Skip to main content

gaze_proxy/
server.rs

1use std::collections::{HashMap, HashSet};
2use std::fmt;
3use std::net::{IpAddr, SocketAddr};
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use axum::body::{Body, Bytes};
8use axum::extract::{ConnectInfo, State};
9use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, Request, StatusCode, Uri};
10use axum::response::{IntoResponse, Response};
11use axum::routing::get;
12use axum::Router;
13use futures_util::StreamExt;
14use gaze::{
15    token_shape, BoundaryRefusal, CleanDocument, CommittedSessionSnapshot, DictionaryBundle,
16    EmittedTokenSpan, LocaleTag, Pipeline, PrefixCacheWriteMode, RawDocument, Scope, Session,
17    SessionTransaction, SessionTransactionError,
18};
19use reqwest::{Client, ClientBuilder};
20use serde::de::{MapAccess, SeqAccess, Visitor};
21use serde::{Deserialize, Deserializer, Serialize};
22use serde_json::{Map, Number, Value};
23use tokio::net::TcpListener;
24use tokio::sync::RwLock;
25use url::{Host, Url};
26
27use crate::adapter::{
28    CoveragePolicy, ProtocolContract, ProviderAdapter, SessionPolicy, SseEvent, SurfaceSyntax,
29};
30use crate::adapters::anthropic::{AnthropicAdapter, DEFAULT_ANTHROPIC_VERSION};
31use crate::codec::{
32    BodyCodec, CodecError, CodecErrorCode, CodecLimits, CodecPhase, InspectionParsedFactsV1,
33    ProvedRequestBody, RequestPseudonymizer, RequestTransformContext, ResponseResidualValidator,
34    ResponseTransformContext, WireFormat,
35};
36use crate::error::{
37    DirectProxyError, HeaderRejectionReason, ProxyError, ProxyErrorCode, ProxyErrorPhase,
38};
39use crate::inspection::{ProxyInspectionEndpointCodesV1, ProxyInspectionLogicalV1};
40use crate::principal::{
41    invoke_resolver, ListenerScope, PeerScope, PrincipalContext, PrincipalResolver,
42    ProcessLocalLoopbackResolver,
43};
44use crate::session_registry::{RegistryError, SessionLease, SessionRegistry};
45use crate::ProxyConfig;
46
47const SESSION_HEADER: &str = "x-gaze-session-id";
48
49/// Compiled ping emitted by a direct SSE response only after its upstream head is accepted.
50pub const DIRECT_PROXY_PING_FRAME: &[u8] = b"event: ping\ndata: {\"type\":\"ping\"}\n\n";
51
52/// The only frame emitted after a direct SSE response has opened and later fails closed.
53pub const DIRECT_PROXY_ERROR_FRAME: &[u8] = b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"proxy_validation_failed\"}}\n\n";
54
55const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
56const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
57const DEFAULT_TOTAL_TIMEOUT: Duration = Duration::from_secs(180);
58const DEFAULT_MAX_REQUEST_BYTES: usize = 2 * 1024 * 1024;
59const DEFAULT_MAX_RESPONSE_BYTES: usize = 32 * 1024 * 1024;
60const DEFAULT_MAX_ERROR_BODY_BYTES: usize = 64 * 1024;
61const DEFAULT_MAX_HEADERS: usize = 64;
62const DEFAULT_MAX_HEADER_NAME_BYTES: usize = 128;
63const DEFAULT_MAX_HEADER_VALUE_BYTES: usize = 8 * 1024;
64const DEFAULT_MAX_HEADER_BYTES: usize = 64 * 1024;
65
66/// Finite direct-client configuration. Additive builders can only lower frozen defaults.
67#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub(crate) struct DirectClientConfig {
69    connect_timeout: Duration,
70    request_timeout: Duration,
71    total_timeout: Duration,
72    max_request_bytes: usize,
73    max_response_bytes: usize,
74    max_error_body_bytes: usize,
75    max_headers: usize,
76    max_header_name_bytes: usize,
77    max_header_value_bytes: usize,
78    max_header_bytes: usize,
79}
80
81impl Default for DirectClientConfig {
82    fn default() -> Self {
83        Self {
84            connect_timeout: DEFAULT_CONNECT_TIMEOUT,
85            request_timeout: DEFAULT_REQUEST_TIMEOUT,
86            total_timeout: DEFAULT_TOTAL_TIMEOUT,
87            max_request_bytes: DEFAULT_MAX_REQUEST_BYTES,
88            max_response_bytes: DEFAULT_MAX_RESPONSE_BYTES,
89            max_error_body_bytes: DEFAULT_MAX_ERROR_BODY_BYTES,
90            max_headers: DEFAULT_MAX_HEADERS,
91            max_header_name_bytes: DEFAULT_MAX_HEADER_NAME_BYTES,
92            max_header_value_bytes: DEFAULT_MAX_HEADER_VALUE_BYTES,
93            max_header_bytes: DEFAULT_MAX_HEADER_BYTES,
94        }
95    }
96}
97
98// P3 builds this crate-private substrate before the P4 orchestration wires it into `serve`.
99#[allow(dead_code)]
100impl DirectClientConfig {
101    fn from_anthropic(adapter: Option<&AnthropicAdapter>) -> Self {
102        let Some(adapter) = adapter else {
103            return Self::default();
104        };
105        let (max_headers, max_header_name_bytes, max_header_value_bytes, max_header_bytes) =
106            adapter.header_limits();
107        Self {
108            connect_timeout: adapter.connect_timeout(),
109            request_timeout: adapter.request_timeout(),
110            total_timeout: adapter.total_timeout(),
111            max_headers,
112            max_header_name_bytes,
113            max_header_value_bytes,
114            max_header_bytes,
115            ..Self::default()
116        }
117    }
118
119    #[must_use]
120    pub(crate) const fn connect_timeout(self) -> Duration {
121        self.connect_timeout
122    }
123
124    #[must_use]
125    pub(crate) const fn request_timeout(self) -> Duration {
126        self.request_timeout
127    }
128
129    #[must_use]
130    pub(crate) const fn total_timeout(self) -> Duration {
131        self.total_timeout
132    }
133
134    #[must_use]
135    pub(crate) const fn max_request_bytes(self) -> usize {
136        self.max_request_bytes
137    }
138
139    #[must_use]
140    pub(crate) const fn max_response_bytes(self) -> usize {
141        self.max_response_bytes
142    }
143
144    #[must_use]
145    pub(crate) const fn max_error_body_bytes(self) -> usize {
146        self.max_error_body_bytes
147    }
148
149    /// Lowers all finite timeouts for a deployment or deterministic test.
150    pub(crate) fn try_with_timeouts(
151        mut self,
152        connect: Duration,
153        request: Duration,
154        total: Duration,
155    ) -> Result<Self, DirectProxyError> {
156        if connect.is_zero()
157            || request.is_zero()
158            || total.is_zero()
159            || connect > DEFAULT_CONNECT_TIMEOUT
160            || request > DEFAULT_REQUEST_TIMEOUT
161            || total > DEFAULT_TOTAL_TIMEOUT
162        {
163            return Err(
164                ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
165            );
166        }
167        self.connect_timeout = connect;
168        self.request_timeout = request;
169        self.total_timeout = total;
170        Ok(self)
171    }
172
173    /// Lowers the successful upstream response-body limit.
174    pub(crate) fn try_with_max_response_bytes(
175        mut self,
176        value: usize,
177    ) -> Result<Self, DirectProxyError> {
178        if value == 0 || value > DEFAULT_MAX_RESPONSE_BYTES {
179            return Err(
180                ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
181            );
182        }
183        self.max_response_bytes = value;
184        Ok(self)
185    }
186
187    /// Lowers the discard-only upstream error-body budget.
188    pub(crate) fn try_with_max_error_body_bytes(
189        mut self,
190        value: usize,
191    ) -> Result<Self, DirectProxyError> {
192        if value == 0 || value > DEFAULT_MAX_ERROR_BODY_BYTES {
193            return Err(
194                ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
195            );
196        }
197        self.max_error_body_bytes = value;
198        Ok(self)
199    }
200
201    /// Lowers the exact outbound request-body limit.
202    pub(crate) fn try_with_max_request_bytes(
203        mut self,
204        value: usize,
205    ) -> Result<Self, DirectProxyError> {
206        if value == 0 || value > DEFAULT_MAX_REQUEST_BYTES {
207            return Err(
208                ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
209            );
210        }
211        self.max_request_bytes = value;
212        Ok(self)
213    }
214}
215
216struct DirectRuntime {
217    client: DirectClient,
218    endpoint: Url,
219    session_policy: SessionPolicy,
220    registry: Option<SessionRegistry>,
221    resolver: Arc<dyn PrincipalResolver>,
222    listener_scope: ListenerScope,
223    allowed_versions: Arc<[String]>,
224    allowed_betas: Arc<[String]>,
225    codec_limits: CodecLimits,
226    client_config: DirectClientConfig,
227    ping_interval: Duration,
228    inspection_endpoint_codes: ProxyInspectionEndpointCodesV1,
229}
230
231enum DirectIngressSession {
232    Ephemeral,
233    Continuity(SessionLease),
234}
235
236struct PreparedDirectIngress {
237    outbound_headers: HeaderMap,
238    session: DirectIngressSession,
239}
240
241impl DirectRuntime {
242    fn from_config(config: &ProxyConfig) -> Result<Option<Self>, ProxyError> {
243        let direct_adapters: Vec<_> = config
244            .adapters
245            .iter()
246            .filter(|adapter| matches!(adapter.contract().protocol(), ProtocolContract::Codec(_)))
247            .collect();
248        if direct_adapters.is_empty() {
249            return Ok(None);
250        }
251        if direct_adapters.len() != 1 {
252            return Err(direct_readiness_error("direct_adapter_count_invalid"));
253        }
254        let adapter = direct_adapters[0];
255        let contract = adapter.contract();
256        let session_policy = contract
257            .session_policy()
258            .ok_or_else(|| direct_readiness_error("direct_session_policy_missing"))?;
259        let codec_limits = contract
260            .codec_limits()
261            .ok_or_else(|| direct_readiness_error("direct_codec_limits_missing"))?;
262        validate_direct_origin(adapter.upstream_base())
263            .map_err(|_| direct_readiness_error("direct_upstream_invalid"))?;
264        let mut endpoint = adapter.upstream_base().clone();
265        endpoint.set_path("/v1/messages");
266        let inspection_endpoint_codes =
267            ProxyInspectionEndpointCodesV1::from_validated_origin(&endpoint);
268
269        let explicit = config.direct_anthropic.as_deref();
270        if explicit.is_some_and(|value| value.upstream_base() != adapter.upstream_base()) {
271            return Err(direct_readiness_error("direct_adapter_mismatch"));
272        }
273        let client_config = DirectClientConfig::from_anthropic(explicit);
274        let client = DirectClient::new(client_config)
275            .map_err(|_| direct_readiness_error("direct_client_invalid"))?;
276        let listener_scope = if config.bind.ip().is_loopback() {
277            ListenerScope::Loopback
278        } else {
279            ListenerScope::NonLoopback
280        };
281        let resolver: Arc<dyn PrincipalResolver> =
282            if let Some(resolver) = explicit.and_then(AnthropicAdapter::principal_resolver) {
283                Arc::clone(resolver)
284            } else if listener_scope == ListenerScope::Loopback {
285                Arc::new(
286                    ProcessLocalLoopbackResolver::new()
287                        .map_err(|_| direct_readiness_error("direct_principal_unavailable"))?,
288                )
289            } else {
290                return Err(direct_readiness_error("direct_principal_required"));
291            };
292        let registry = match session_policy {
293            SessionPolicy::EphemeralSingleRequest => None,
294            SessionPolicy::OpaqueHeaderContinuity(registry_config) => Some(
295                SessionRegistry::new(registry_config)
296                    .map_err(|_| direct_readiness_error("direct_registry_invalid"))?,
297            ),
298        };
299        let allowed_versions = explicit.map_or_else(
300            || Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
301            |value| Arc::from(value.allowed_versions()),
302        );
303        let allowed_betas = explicit.map_or_else(
304            || Arc::<[String]>::from([]),
305            |value| Arc::from(value.allowed_betas()),
306        );
307        let ping_interval =
308            explicit.map_or(Duration::from_secs(10), AnthropicAdapter::ping_interval);
309        Ok(Some(Self {
310            client,
311            endpoint,
312            session_policy,
313            registry,
314            resolver,
315            listener_scope,
316            allowed_versions,
317            allowed_betas,
318            codec_limits,
319            client_config,
320            ping_interval,
321            inspection_endpoint_codes,
322        }))
323    }
324}
325
326fn direct_readiness_error(detail: &'static str) -> ProxyError {
327    ProxyError::DaemonConfig {
328        detail: detail.to_string(),
329    }
330}
331
332/// One exact already-proved direct request. It has no `Debug` implementation by design.
333#[allow(dead_code)]
334pub(crate) struct DirectRequest {
335    url: Url,
336    headers: HeaderMap,
337    body: Vec<u8>,
338    expected_response: WireFormat,
339    inspection_parsed_facts: Option<InspectionParsedFactsV1>,
340    signed_or_encrypted_surface: bool,
341}
342
343#[allow(dead_code)]
344impl DirectRequest {
345    /// Consumes the exact codec-proved body after validating the safe upstream URL and headers.
346    pub(crate) fn from_proved(
347        url: Url,
348        headers: HeaderMap,
349        proved_body: ProvedRequestBody,
350        expected_response: WireFormat,
351    ) -> Result<Self, DirectProxyError> {
352        validate_direct_url(&url)?;
353        validate_outbound_headers(&headers)?;
354        if !proved_body.provenance().final_buffer_verified() {
355            return Err(ProxyErrorCode::InvalidProvenance.error(ProxyErrorPhase::RequestTransform));
356        }
357        let (body, inspection_parsed_facts, signed_or_encrypted_surface) =
358            proved_body.into_inspection_parts();
359        Ok(Self {
360            url,
361            headers,
362            body,
363            expected_response,
364            inspection_parsed_facts,
365            signed_or_encrypted_surface,
366        })
367    }
368}
369
370/// A proved request whose real core transaction has not yet been published.
371///
372/// This type is crate-private and has exactly one consuming transition to
373/// [`CommittedDirectRequest`]. Dropping it publishes neither mappings nor prefix-cache state.
374struct PreparedDirectRequest<'session> {
375    transaction: SessionTransaction<'session>,
376    request: DirectRequest,
377    lifecycle: DirectTransactionState,
378}
379
380impl<'session> PreparedDirectRequest<'session> {
381    fn new(transaction: SessionTransaction<'session>, request: DirectRequest) -> Self {
382        Self {
383            transaction,
384            request,
385            lifecycle: DirectTransactionState::new(),
386        }
387    }
388
389    fn commit(mut self) -> Result<CommittedDirectRequest, SessionTransactionError> {
390        let snapshot = self.transaction.commit()?;
391        self.lifecycle.state = DirectIoState::CommittedNoIo;
392        Ok(CommittedDirectRequest {
393            snapshot,
394            request: self.request,
395            lifecycle: self.lifecycle,
396        })
397    }
398}
399
400/// The only request state accepted by [`DirectClient`].
401///
402/// Fields stay private so neither the committed snapshot nor the no-I/O state can be forged.
403struct CommittedDirectRequest {
404    snapshot: CommittedSessionSnapshot,
405    request: DirectRequest,
406    lifecycle: DirectTransactionState,
407}
408
409impl fmt::Debug for CommittedDirectRequest {
410    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
411        formatter
412            .debug_struct("CommittedDirectRequest")
413            .field("lifecycle", &self.lifecycle)
414            .finish_non_exhaustive()
415    }
416}
417
418struct CommittedDirectResponse {
419    snapshot: CommittedSessionSnapshot,
420    response: DirectResponse,
421    lifecycle: DirectTransactionState,
422}
423
424struct CommittedDirectSse {
425    snapshot: CommittedSessionSnapshot,
426    response: reqwest::Response,
427    head: DirectResponseHead,
428    lifecycle: DirectTransactionState,
429    remaining_total_timeout: Duration,
430}
431
432impl fmt::Debug for CommittedDirectResponse {
433    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
434        formatter
435            .debug_struct("CommittedDirectResponse")
436            .field("status", &self.response.status())
437            .field("format", &self.response.format())
438            .field("lifecycle", &self.lifecycle)
439            .finish_non_exhaustive()
440    }
441}
442
443/// Validated minimal upstream response head.
444pub struct DirectResponseHead {
445    status: StatusCode,
446    format: WireFormat,
447    headers: HeaderMap,
448}
449
450impl fmt::Debug for DirectResponseHead {
451    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
452        formatter
453            .debug_struct("DirectResponseHead")
454            .field("status", &self.status)
455            .field("format", &self.format)
456            .field("header_count", &self.headers.len())
457            .finish()
458    }
459}
460
461impl DirectResponseHead {
462    #[must_use]
463    pub const fn status(&self) -> StatusCode {
464        self.status
465    }
466
467    #[must_use]
468    pub const fn format(&self) -> WireFormat {
469        self.format
470    }
471
472    #[must_use]
473    pub const fn headers(&self) -> &HeaderMap {
474        &self.headers
475    }
476}
477
478/// A buffered direct response whose head and accumulated body have both passed substrate checks.
479#[allow(dead_code)]
480pub(crate) struct DirectResponse {
481    head: DirectResponseHead,
482    body: Vec<u8>,
483}
484
485impl fmt::Debug for DirectResponse {
486    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
487        formatter
488            .debug_struct("DirectResponse")
489            .field("status", &self.head.status)
490            .field("format", &self.head.format)
491            .field("body_bytes", &self.body.len())
492            .finish()
493    }
494}
495
496#[allow(dead_code)]
497impl DirectResponse {
498    #[must_use]
499    pub(crate) const fn status(&self) -> StatusCode {
500        self.head.status()
501    }
502
503    #[must_use]
504    pub(crate) const fn format(&self) -> WireFormat {
505        self.head.format()
506    }
507
508    #[must_use]
509    pub(crate) const fn headers(&self) -> &HeaderMap {
510        self.head.headers()
511    }
512
513    #[must_use]
514    pub(crate) fn body(&self) -> &[u8] {
515        &self.body
516    }
517
518    #[must_use]
519    pub(crate) fn into_body(self) -> Vec<u8> {
520        self.body
521    }
522
523    fn into_parts(self) -> (DirectResponseHead, Vec<u8>) {
524        (self.head, self.body)
525    }
526}
527
528/// Dedicated hardened client for direct codec-backed adapters.
529#[derive(Clone)]
530#[allow(dead_code)]
531pub(crate) struct DirectClient {
532    client: Client,
533    config: DirectClientConfig,
534}
535
536#[allow(dead_code)]
537impl DirectClient {
538    /// Builds a client with redirects, proxies, retries, referrers, and decompression disabled.
539    pub(crate) fn new(config: DirectClientConfig) -> Result<Self, DirectProxyError> {
540        let client = hardened_client_builder(config).build().map_err(|_| {
541            ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::UpstreamConfiguration)
542        })?;
543        Ok(Self { client, config })
544    }
545
546    /// Sends the exact request body once and buffers the raw response with accumulated limits.
547    async fn execute(
548        &self,
549        mut committed: CommittedDirectRequest,
550    ) -> Result<CommittedDirectResponse, DirectProxyError> {
551        committed.lifecycle.require_committed_no_io()?;
552        if committed.request.body.len() > self.config.max_request_bytes {
553            return Err(
554                ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
555            );
556        }
557        committed.lifecycle.mark_io_attempted()?;
558        match tokio::time::timeout(
559            self.config.total_timeout,
560            self.execute_inner(committed.request),
561        )
562        .await
563        {
564            Ok(Ok(response)) => Ok(CommittedDirectResponse {
565                snapshot: committed.snapshot,
566                response,
567                lifecycle: committed.lifecycle,
568            }),
569            Ok(Err(error)) => Err(error),
570            Err(_) => Err(ProxyErrorCode::TotalTimeout.error(ProxyErrorPhase::UpstreamBody)),
571        }
572    }
573
574    async fn execute_sse(
575        &self,
576        mut committed: CommittedDirectRequest,
577    ) -> Result<CommittedDirectSse, DirectProxyError> {
578        committed.lifecycle.require_committed_no_io()?;
579        if committed.request.body.len() > self.config.max_request_bytes {
580            return Err(
581                ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
582            );
583        }
584        if committed.request.expected_response != WireFormat::Sse {
585            return Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework));
586        }
587        committed.lifecycle.mark_io_attempted()?;
588        let started = tokio::time::Instant::now();
589        let opened = tokio::time::timeout(
590            self.config.total_timeout,
591            self.open_sse_response(committed.request),
592        )
593        .await
594        .map_err(|_| ProxyErrorCode::TotalTimeout.error(ProxyErrorPhase::UpstreamHeaders))??;
595        let remaining_total_timeout = self.config.total_timeout.saturating_sub(started.elapsed());
596        Ok(CommittedDirectSse {
597            snapshot: committed.snapshot,
598            response: opened.0,
599            head: opened.1,
600            lifecycle: committed.lifecycle,
601            remaining_total_timeout,
602        })
603    }
604
605    async fn open_sse_response(
606        &self,
607        request: DirectRequest,
608    ) -> Result<(reqwest::Response, DirectResponseHead), DirectProxyError> {
609        let response = self
610            .client
611            .post(request.url)
612            .headers(request.headers)
613            .header(reqwest::header::ACCEPT_ENCODING, "identity")
614            .body(request.body)
615            .timeout(self.config.request_timeout)
616            .send()
617            .await
618            .map_err(map_reqwest_error)?;
619        let status = response.status();
620        if status.is_redirection() {
621            return Err(ProxyErrorCode::UpstreamRedirect.error(ProxyErrorPhase::UpstreamHeaders));
622        }
623        match validate_direct_response_head_with_config(
624            status,
625            response.headers(),
626            WireFormat::Sse,
627            self.config,
628        ) {
629            Ok(head) => Ok((response, head)),
630            Err(error) if is_mapped_upstream_status(error.code()) => {
631                discard_error_body(response, self.config.max_error_body_bytes).await;
632                Err(error)
633            }
634            Err(error) => Err(error),
635        }
636    }
637
638    async fn execute_inner(
639        &self,
640        request: DirectRequest,
641    ) -> Result<DirectResponse, DirectProxyError> {
642        let mut outbound = self
643            .client
644            .post(request.url)
645            .headers(request.headers)
646            .header(reqwest::header::ACCEPT_ENCODING, "identity")
647            .body(request.body);
648        outbound = outbound.timeout(self.config.request_timeout);
649        let response = outbound.send().await.map_err(map_reqwest_error)?;
650        let status = response.status();
651
652        if status.is_redirection() {
653            return Err(ProxyErrorCode::UpstreamRedirect.error(ProxyErrorPhase::UpstreamHeaders));
654        }
655
656        let head = validate_direct_response_head_with_config(
657            status,
658            response.headers(),
659            request.expected_response,
660            self.config,
661        );
662        match head {
663            Ok(head) => {
664                let body = collect_response_body(response, self.config.max_response_bytes).await?;
665                validate_declared_body(
666                    request.expected_response,
667                    &body,
668                    self.config.max_response_bytes,
669                )?;
670                Ok(DirectResponse { head, body })
671            }
672            Err(error) if is_mapped_upstream_status(error.code()) => {
673                discard_error_body(response, self.config.max_error_body_bytes).await;
674                Err(error)
675            }
676            Err(error) => Err(error),
677        }
678    }
679}
680
681#[allow(dead_code)]
682fn hardened_client_builder(config: DirectClientConfig) -> ClientBuilder {
683    Client::builder()
684        .redirect(reqwest::redirect::Policy::none())
685        .no_proxy()
686        .referer(false)
687        .retry(reqwest::retry::never())
688        .no_gzip()
689        .no_brotli()
690        .no_zstd()
691        .no_deflate()
692        .connect_timeout(config.connect_timeout)
693        .timeout(config.request_timeout)
694}
695
696#[allow(dead_code)]
697fn map_reqwest_error(error: reqwest::Error) -> DirectProxyError {
698    let code = if error.is_timeout() && error.is_connect() {
699        ProxyErrorCode::ConnectTimeout
700    } else if error.is_timeout() {
701        ProxyErrorCode::RequestTimeout
702    } else if error.is_connect() {
703        ProxyErrorCode::UpstreamUnreachable
704    } else {
705        ProxyErrorCode::UpstreamProtocol
706    };
707    code.error(ProxyErrorPhase::UpstreamConnect)
708}
709
710#[allow(dead_code)]
711fn map_response_body_error(error: reqwest::Error) -> DirectProxyError {
712    let code = if error.is_timeout() {
713        ProxyErrorCode::RequestTimeout
714    } else {
715        ProxyErrorCode::UpstreamProtocol
716    };
717    code.error(ProxyErrorPhase::UpstreamBody)
718}
719
720#[allow(dead_code)]
721async fn collect_response_body(
722    response: reqwest::Response,
723    limit: usize,
724) -> Result<Vec<u8>, DirectProxyError> {
725    let mut stream = response.bytes_stream();
726    let mut body = Vec::new();
727    while let Some(chunk) = stream.next().await {
728        let chunk = chunk.map_err(map_response_body_error)?;
729        let accumulated = body.len().checked_add(chunk.len()).ok_or_else(|| {
730            ProxyErrorCode::ResponseBodyLimitExceeded.error(ProxyErrorPhase::UpstreamBody)
731        })?;
732        if accumulated > limit {
733            return Err(
734                ProxyErrorCode::ResponseBodyLimitExceeded.error(ProxyErrorPhase::UpstreamBody)
735            );
736        }
737        body.extend_from_slice(&chunk);
738    }
739    Ok(body)
740}
741
742#[allow(dead_code)]
743async fn discard_error_body(response: reqwest::Response, limit: usize) {
744    let mut stream = response.bytes_stream();
745    let mut discarded = 0_usize;
746    while discarded < limit {
747        let Some(next) = stream.next().await else {
748            break;
749        };
750        let Ok(chunk) = next else {
751            break;
752        };
753        let remaining = limit - discarded;
754        discarded += chunk.len().min(remaining);
755        if chunk.len() > remaining {
756            break;
757        }
758    }
759}
760
761#[allow(dead_code)]
762fn validate_direct_url(url: &Url) -> Result<(), DirectProxyError> {
763    let safe_authority = url.username().is_empty()
764        && url.password().is_none()
765        && url.query().is_none()
766        && url.fragment().is_none()
767        && url.host().is_some();
768    let safe_scheme = match (url.scheme(), url.host()) {
769        ("https", Some(_)) => true,
770        ("http", Some(Host::Ipv4(address))) => IpAddr::V4(address).is_loopback(),
771        ("http", Some(Host::Ipv6(address))) => IpAddr::V6(address).is_loopback(),
772        _ => false,
773    };
774    if safe_authority && safe_scheme {
775        Ok(())
776    } else {
777        Err(ProxyErrorCode::InvalidUpstreamUrl.error(ProxyErrorPhase::UpstreamConfiguration))
778    }
779}
780
781fn validate_direct_origin(url: &Url) -> Result<(), DirectProxyError> {
782    if url.path() != "/" {
783        return Err(
784            ProxyErrorCode::InvalidUpstreamUrl.error(ProxyErrorPhase::UpstreamConfiguration)
785        );
786    }
787    validate_direct_url(url)
788}
789
790#[allow(dead_code)]
791fn validate_outbound_headers(headers: &HeaderMap) -> Result<(), DirectProxyError> {
792    const ALLOWED: &[&str] = &[
793        "content-type",
794        "x-api-key",
795        "anthropic-version",
796        "anthropic-beta",
797    ];
798    validate_header_space(headers, DirectClientConfig::default())?;
799    let mut singleton_names = HashSet::new();
800    for (name, value) in headers {
801        if !ALLOWED.contains(&name.as_str()) {
802            return Err(ProxyErrorCode::HeaderRejected
803                .error(ProxyErrorPhase::RequestValidation)
804                .with_header_rejection(name.as_str(), HeaderRejectionReason::NameNotAllowlisted));
805        }
806        if value.is_empty() {
807            return Err(ProxyErrorCode::HeaderRejected
808                .error(ProxyErrorPhase::RequestValidation)
809                .with_header_rejection(name.as_str(), HeaderRejectionReason::EmptyValue));
810        }
811        if matches!(
812            name.as_str(),
813            "x-api-key" | "anthropic-version" | "anthropic-beta"
814        ) && !singleton_names.insert(name.as_str())
815        {
816            return Err(ProxyErrorCode::HeaderRejected
817                .error(ProxyErrorPhase::RequestValidation)
818                .with_header_rejection(name.as_str(), HeaderRejectionReason::DuplicateSingleton));
819        }
820    }
821    Ok(())
822}
823
824/// Maps an upstream status without reading its body or redirect metadata.
825pub fn classify_upstream_status(status: StatusCode) -> Result<(), DirectProxyError> {
826    let code = if status.is_success() {
827        return Ok(());
828    } else if status.is_redirection() {
829        ProxyErrorCode::UpstreamRedirect
830    } else {
831        match status.as_u16() {
832            400 => ProxyErrorCode::UpstreamBadRequest,
833            401 => ProxyErrorCode::UpstreamUnauthorized,
834            403 => ProxyErrorCode::UpstreamForbidden,
835            404 => ProxyErrorCode::UpstreamNotFound,
836            409 => ProxyErrorCode::UpstreamConflict,
837            413 => ProxyErrorCode::UpstreamPayloadTooLarge,
838            429 => ProxyErrorCode::UpstreamRateLimited,
839            502 | 503 | 504 | 529 => ProxyErrorCode::UpstreamUnavailable,
840            value if (400..=499).contains(&value) => ProxyErrorCode::UpstreamClientFailure,
841            500..=599 => ProxyErrorCode::UpstreamServerFailure,
842            _ => ProxyErrorCode::UpstreamServerFailure,
843        }
844    };
845    Err(code.error(ProxyErrorPhase::UpstreamHeaders))
846}
847
848#[allow(dead_code)]
849fn is_mapped_upstream_status(code: ProxyErrorCode) -> bool {
850    matches!(
851        code,
852        ProxyErrorCode::UpstreamBadRequest
853            | ProxyErrorCode::UpstreamUnauthorized
854            | ProxyErrorCode::UpstreamForbidden
855            | ProxyErrorCode::UpstreamNotFound
856            | ProxyErrorCode::UpstreamConflict
857            | ProxyErrorCode::UpstreamPayloadTooLarge
858            | ProxyErrorCode::UpstreamRateLimited
859            | ProxyErrorCode::UpstreamClientFailure
860            | ProxyErrorCode::UpstreamUnavailable
861            | ProxyErrorCode::UpstreamServerFailure
862    )
863}
864
865/// Validates a raw upstream head and rebuilds its minimal transformed header allowlist.
866///
867/// Header limits, framing, and retry metadata are validated before an error status is mapped, so
868/// malformed metadata can never be hidden behind a provider status code.
869pub fn validate_direct_response_head(
870    status: StatusCode,
871    headers: &HeaderMap,
872    expected: WireFormat,
873) -> Result<DirectResponseHead, DirectProxyError> {
874    validate_direct_response_head_with_config(
875        status,
876        headers,
877        expected,
878        DirectClientConfig::default(),
879    )
880}
881
882fn validate_direct_response_head_with_config(
883    status: StatusCode,
884    headers: &HeaderMap,
885    expected: WireFormat,
886    config: DirectClientConfig,
887) -> Result<DirectResponseHead, DirectProxyError> {
888    if status.is_redirection() {
889        return Err(ProxyErrorCode::UpstreamRedirect.error(ProxyErrorPhase::UpstreamHeaders));
890    }
891    validate_header_space(headers, config)?;
892    validate_response_framing(headers)?;
893
894    if let Err(error) = classify_upstream_status(status) {
895        let retry_after_seconds = parse_retry_after(headers)?;
896        return Err(error.with_retry_after_seconds(retry_after_seconds));
897    }
898
899    let content_type = one_header(headers, reqwest::header::CONTENT_TYPE)?;
900    let actual = match content_type {
901        None if expected == WireFormat::Empty => WireFormat::Empty,
902        Some(value) => parse_wire_content_type(value)?,
903        None => {
904            return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
905                .error(ProxyErrorPhase::ResponseValidation));
906        }
907    };
908    if actual != expected || (expected == WireFormat::Empty && content_type.is_some()) {
909        return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
910            .error(ProxyErrorPhase::ResponseValidation));
911    }
912    let mut minimal = HeaderMap::new();
913    if expected != WireFormat::Empty {
914        minimal.insert(
915            reqwest::header::CONTENT_TYPE,
916            HeaderValue::from_static(canonical_content_type(expected)),
917        );
918    }
919    Ok(DirectResponseHead {
920        status,
921        format: expected,
922        headers: minimal,
923    })
924}
925
926/// Validates one complete declared body without format sniffing or fallback.
927pub fn validate_declared_body(
928    format: WireFormat,
929    body: &[u8],
930    limit: usize,
931) -> Result<(), DirectProxyError> {
932    if body.len() > limit {
933        return Err(
934            ProxyErrorCode::ResponseBodyLimitExceeded.error(ProxyErrorPhase::ResponseValidation)
935        );
936    }
937    match format {
938        WireFormat::Empty if body.is_empty() => Ok(()),
939        WireFormat::Empty => Err(ProxyErrorCode::InvalidUpstreamResponseFormat
940            .error(ProxyErrorPhase::ResponseValidation)),
941        WireFormat::Json => validate_unique_json(body),
942        WireFormat::Ndjson => {
943            let text = std::str::from_utf8(body).map_err(|_| {
944                ProxyErrorCode::InvalidUpstreamResponseFormat
945                    .error(ProxyErrorPhase::ResponseValidation)
946            })?;
947            for line in text.lines().map(str::trim).filter(|line| !line.is_empty()) {
948                validate_unique_json(line.as_bytes())?;
949            }
950            Ok(())
951        }
952        WireFormat::Utf8Text => std::str::from_utf8(body).map(|_| ()).map_err(|_| {
953            ProxyErrorCode::InvalidUpstreamResponseFormat.error(ProxyErrorPhase::ResponseValidation)
954        }),
955        WireFormat::Sse => validate_sse_envelope(body),
956    }
957}
958
959fn validate_unique_json(body: &[u8]) -> Result<(), DirectProxyError> {
960    let parsed: UniqueJson = serde_json::from_slice(body).map_err(|_| {
961        ProxyErrorCode::InvalidUpstreamResponseFormat.error(ProxyErrorPhase::ResponseValidation)
962    })?;
963    if parsed.duplicate {
964        Err(ProxyErrorCode::DuplicateObjectKey.error(ProxyErrorPhase::ResponseValidation))
965    } else {
966        Ok(())
967    }
968}
969
970fn validate_sse_envelope(body: &[u8]) -> Result<(), DirectProxyError> {
971    if body.starts_with(&[0xef, 0xbb, 0xbf]) || body.contains(&0) {
972        return Err(ProxyErrorCode::InvalidSseLifecycle.error(ProxyErrorPhase::ResponseValidation));
973    }
974    std::str::from_utf8(body).map_err(|_| {
975        ProxyErrorCode::InvalidUpstreamResponseFormat.error(ProxyErrorPhase::ResponseValidation)
976    })?;
977    if body.ends_with(b"\n\n") || body.ends_with(b"\r\r") || body.ends_with(b"\r\n\r\n") {
978        Ok(())
979    } else {
980        Err(ProxyErrorCode::InvalidSseLifecycle.error(ProxyErrorPhase::ResponseValidation))
981    }
982}
983
984fn validate_header_space(
985    headers: &HeaderMap,
986    config: DirectClientConfig,
987) -> Result<(), DirectProxyError> {
988    let mut count = 0_usize;
989    let mut aggregate = 0_usize;
990    for (name, value) in headers {
991        count = count.checked_add(1).ok_or_else(header_limit_error)?;
992        aggregate = aggregate
993            .checked_add(name.as_str().len())
994            .and_then(|value_so_far| value_so_far.checked_add(value.as_bytes().len()))
995            .ok_or_else(header_limit_error)?;
996        if count > config.max_headers
997            || name.as_str().len() > config.max_header_name_bytes
998            || value.as_bytes().len() > config.max_header_value_bytes
999            || aggregate > config.max_header_bytes
1000        {
1001            return Err(header_limit_error());
1002        }
1003    }
1004    Ok(())
1005}
1006
1007fn header_limit_error() -> DirectProxyError {
1008    ProxyErrorCode::HeaderLimitExceeded.error(ProxyErrorPhase::UpstreamHeaders)
1009}
1010
1011fn one_header(
1012    headers: &HeaderMap,
1013    name: HeaderName,
1014) -> Result<Option<&HeaderValue>, DirectProxyError> {
1015    let mut values = headers.get_all(name).iter();
1016    let first = values.next();
1017    if values.next().is_some() {
1018        return Err(ProxyErrorCode::InvalidUpstreamHeader.error(ProxyErrorPhase::UpstreamHeaders));
1019    }
1020    Ok(first)
1021}
1022
1023fn validate_response_framing(headers: &HeaderMap) -> Result<(), DirectProxyError> {
1024    let content_length = one_header(headers, reqwest::header::CONTENT_LENGTH)?;
1025    let transfer_encoding = one_header(headers, reqwest::header::TRANSFER_ENCODING)?;
1026    if content_length.is_some() && transfer_encoding.is_some() {
1027        return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::UpstreamHeaders));
1028    }
1029    if let Some(value) = content_length {
1030        let raw = value.as_bytes();
1031        if raw.is_empty() || !raw.iter().all(u8::is_ascii_digit) {
1032            return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::UpstreamHeaders));
1033        }
1034    }
1035    if let Some(value) = transfer_encoding {
1036        if !value.as_bytes().eq_ignore_ascii_case(b"chunked") {
1037            return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::UpstreamHeaders));
1038        }
1039    }
1040    if let Some(value) = one_header(headers, reqwest::header::CONTENT_ENCODING)? {
1041        if !trim_ascii(value.as_bytes()).eq_ignore_ascii_case(b"identity") {
1042            return Err(
1043                ProxyErrorCode::UnsupportedContentEncoding.error(ProxyErrorPhase::UpstreamHeaders)
1044            );
1045        }
1046    }
1047    Ok(())
1048}
1049
1050fn parse_retry_after(headers: &HeaderMap) -> Result<Option<u32>, DirectProxyError> {
1051    let Some(value) = one_header(headers, reqwest::header::RETRY_AFTER)? else {
1052        return Ok(None);
1053    };
1054    let bytes = value.as_bytes();
1055    if bytes.is_empty() || !bytes.iter().all(u8::is_ascii_digit) {
1056        return Err(ProxyErrorCode::InvalidUpstreamHeader.error(ProxyErrorPhase::UpstreamHeaders));
1057    }
1058    let seconds = std::str::from_utf8(bytes)
1059        .ok()
1060        .and_then(|text| text.parse::<u32>().ok())
1061        .ok_or_else(|| {
1062            ProxyErrorCode::InvalidUpstreamHeader.error(ProxyErrorPhase::UpstreamHeaders)
1063        })?;
1064    Ok(Some(seconds))
1065}
1066
1067fn parse_wire_content_type(value: &HeaderValue) -> Result<WireFormat, DirectProxyError> {
1068    let raw = value.as_bytes();
1069    if !raw.is_ascii() {
1070        return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1071            .error(ProxyErrorPhase::ResponseValidation));
1072    }
1073    let mut parts = raw.split(|byte| *byte == b';');
1074    let media_type = trim_ascii(parts.next().unwrap_or_default());
1075    let format = if media_type.eq_ignore_ascii_case(b"application/json") {
1076        WireFormat::Json
1077    } else if media_type.eq_ignore_ascii_case(b"application/x-ndjson") {
1078        WireFormat::Ndjson
1079    } else if media_type.eq_ignore_ascii_case(b"text/plain") {
1080        WireFormat::Utf8Text
1081    } else if media_type.eq_ignore_ascii_case(b"text/event-stream") {
1082        WireFormat::Sse
1083    } else {
1084        return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1085            .error(ProxyErrorPhase::ResponseValidation));
1086    };
1087    for parameter in parts {
1088        validate_mime_parameter(trim_ascii(parameter))?;
1089    }
1090    Ok(format)
1091}
1092
1093fn validate_mime_parameter(parameter: &[u8]) -> Result<(), DirectProxyError> {
1094    let Some(separator) = parameter.iter().position(|byte| *byte == b'=') else {
1095        return Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1096            .error(ProxyErrorPhase::ResponseValidation));
1097    };
1098    let name = trim_ascii(&parameter[..separator]);
1099    let value = trim_ascii(&parameter[separator + 1..]);
1100    let valid_name = !name.is_empty() && name.iter().all(|byte| is_mime_token(*byte));
1101    let valid_value = if value.len() >= 2 && value[0] == b'"' && value[value.len() - 1] == b'"' {
1102        value[1..value.len() - 1]
1103            .iter()
1104            .all(|byte| matches!(*byte, b'\t' | 0x20..=0x7e) && *byte != b'"' && *byte != b'\\')
1105    } else {
1106        !value.is_empty() && value.iter().all(|byte| is_mime_token(*byte))
1107    };
1108    if valid_name && valid_value {
1109        Ok(())
1110    } else {
1111        Err(ProxyErrorCode::InvalidUpstreamResponseFormat
1112            .error(ProxyErrorPhase::ResponseValidation))
1113    }
1114}
1115
1116fn is_mime_token(byte: u8) -> bool {
1117    byte.is_ascii_alphanumeric()
1118        || matches!(
1119            byte,
1120            b'!' | b'#'
1121                | b'$'
1122                | b'%'
1123                | b'&'
1124                | b'\''
1125                | b'*'
1126                | b'+'
1127                | b'-'
1128                | b'.'
1129                | b'^'
1130                | b'_'
1131                | b'`'
1132                | b'|'
1133                | b'~'
1134        )
1135}
1136
1137fn trim_ascii(mut value: &[u8]) -> &[u8] {
1138    while matches!(value.first(), Some(b' ' | b'\t')) {
1139        value = &value[1..];
1140    }
1141    while matches!(value.last(), Some(b' ' | b'\t')) {
1142        value = &value[..value.len() - 1];
1143    }
1144    value
1145}
1146
1147fn canonical_content_type(format: WireFormat) -> &'static str {
1148    match format {
1149        WireFormat::Empty => "application/octet-stream",
1150        WireFormat::Json => "application/json",
1151        WireFormat::Ndjson => "application/x-ndjson",
1152        WireFormat::Utf8Text => "text/plain; charset=utf-8",
1153        WireFormat::Sse => "text/event-stream",
1154    }
1155}
1156
1157struct UniqueJson {
1158    value: Value,
1159    duplicate: bool,
1160}
1161
1162impl<'de> Deserialize<'de> for UniqueJson {
1163    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1164    where
1165        D: Deserializer<'de>,
1166    {
1167        deserializer.deserialize_any(UniqueJsonVisitor)
1168    }
1169}
1170
1171struct UniqueJsonVisitor;
1172
1173impl<'de> Visitor<'de> for UniqueJsonVisitor {
1174    type Value = UniqueJson;
1175
1176    fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1177        formatter.write_str("one bounded JSON value")
1178    }
1179
1180    fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
1181        Ok(unique(Value::Bool(value)))
1182    }
1183
1184    fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
1185        Ok(unique(Value::Number(Number::from(value))))
1186    }
1187
1188    fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
1189        Ok(unique(Value::Number(Number::from(value))))
1190    }
1191
1192    fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
1193    where
1194        E: serde::de::Error,
1195    {
1196        Number::from_f64(value)
1197            .map(Value::Number)
1198            .map(unique)
1199            .ok_or_else(|| E::custom("non-finite number"))
1200    }
1201
1202    fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
1203        Ok(unique(Value::String(value.to_owned())))
1204    }
1205
1206    fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
1207        Ok(unique(Value::String(value)))
1208    }
1209
1210    fn visit_none<E>(self) -> Result<Self::Value, E> {
1211        Ok(unique(Value::Null))
1212    }
1213
1214    fn visit_unit<E>(self) -> Result<Self::Value, E> {
1215        Ok(unique(Value::Null))
1216    }
1217
1218    fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
1219    where
1220        A: SeqAccess<'de>,
1221    {
1222        let mut values = Vec::new();
1223        let mut duplicate = false;
1224        while let Some(value) = sequence.next_element::<UniqueJson>()? {
1225            duplicate |= value.duplicate;
1226            values.push(value.value);
1227        }
1228        Ok(UniqueJson {
1229            value: Value::Array(values),
1230            duplicate,
1231        })
1232    }
1233
1234    fn visit_map<A>(self, mut object: A) -> Result<Self::Value, A::Error>
1235    where
1236        A: MapAccess<'de>,
1237    {
1238        let mut values = Map::new();
1239        let mut seen = HashSet::new();
1240        let mut duplicate = false;
1241        while let Some(key) = object.next_key::<String>()? {
1242            duplicate |= !seen.insert(key.clone());
1243            let value = object.next_value::<UniqueJson>()?;
1244            duplicate |= value.duplicate;
1245            values.insert(key, value.value);
1246        }
1247        Ok(UniqueJson {
1248            value: Value::Object(values),
1249            duplicate,
1250        })
1251    }
1252}
1253
1254fn unique(value: Value) -> UniqueJson {
1255    UniqueJson {
1256        value,
1257        duplicate: false,
1258    }
1259}
1260
1261/// Parses the request `stream` control as a closed boolean without sniffing or fallback.
1262pub fn parse_request_stream_format(body: &[u8]) -> Result<WireFormat, DirectProxyError> {
1263    let parsed: UniqueJson = serde_json::from_slice(body).map_err(|_| {
1264        ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation)
1265    })?;
1266    if parsed.duplicate {
1267        return Err(ProxyErrorCode::DuplicateObjectKey.error(ProxyErrorPhase::RequestValidation));
1268    }
1269    let Value::Object(object) = parsed.value else {
1270        return Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation));
1271    };
1272    match object.get("stream") {
1273        None | Some(Value::Bool(false)) => Ok(WireFormat::Json),
1274        Some(Value::Bool(true)) => Ok(WireFormat::Sse),
1275        _ => Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation)),
1276    }
1277}
1278
1279/// Production request protection over the real staged core transaction.
1280///
1281/// It delegates to the single shared Pipeline algorithm introduced by P4a and performs only
1282/// validation against the staged token namespace. No throwaway session, replay, or whole-value
1283/// token shortcut exists here.
1284struct PipelineRequestPseudonymizer<'pipeline, 'context, 'transaction, 'session> {
1285    pipeline: &'pipeline Pipeline,
1286    transaction: &'transaction mut SessionTransaction<'session>,
1287    locale_chain: &'context [LocaleTag],
1288    dictionaries: &'context DictionaryBundle,
1289    /// Why protection refused, kept aside because the codec carries only a closed code.
1290    refusal: Option<BoundaryRefusal>,
1291}
1292
1293impl<'pipeline, 'context, 'transaction, 'session>
1294    PipelineRequestPseudonymizer<'pipeline, 'context, 'transaction, 'session>
1295{
1296    fn new(
1297        pipeline: &'pipeline Pipeline,
1298        locale_chain: &'context [LocaleTag],
1299        dictionaries: &'context DictionaryBundle,
1300        transaction: &'transaction mut SessionTransaction<'session>,
1301    ) -> Self {
1302        Self {
1303            pipeline,
1304            transaction,
1305            locale_chain,
1306            dictionaries,
1307            refusal: None,
1308        }
1309    }
1310
1311    /// Both prefix-cache modes rescan the complete input, so the mode changes nothing here.
1312    fn protect_unstaged_segment(
1313        &mut self,
1314        input: &str,
1315        _prefix_cache_write_mode: PrefixCacheWriteMode,
1316    ) -> Result<String, CodecErrorCode> {
1317        // `gaze clean --safety-net-fallback strict`'s Resolve step; admission
1318        // follows on the leaf.
1319        let result = self.pipeline.resolve_boundary_text_transaction(
1320            self.transaction,
1321            input,
1322            self.locale_chain,
1323            self.dictionaries,
1324        );
1325        self.refused(result)
1326    }
1327
1328    fn refused<T>(&mut self, result: Result<T, BoundaryRefusal>) -> Result<T, CodecErrorCode> {
1329        result.map_err(|refusal| {
1330            self.refusal = Some(refusal);
1331            CodecErrorCode::ProtectionFailedClosed
1332        })
1333    }
1334}
1335
1336impl RequestPseudonymizer for PipelineRequestPseudonymizer<'_, '_, '_, '_> {
1337    fn protect(&mut self, input: &str) -> Result<String, CodecErrorCode> {
1338        self.protect_with_prefix_cache_write_mode(input, PrefixCacheWriteMode::Allow)
1339    }
1340
1341    fn protect_with_prefix_cache_write_mode(
1342        &mut self,
1343        input: &str,
1344        prefix_cache_write_mode: PrefixCacheWriteMode,
1345    ) -> Result<String, CodecErrorCode> {
1346        let output = self.protect_preserving_tokens(input, prefix_cache_write_mode)?;
1347        let admitted = self.pipeline.admit_boundary_text_transaction(
1348            self.transaction,
1349            &output,
1350            self.locale_chain,
1351            self.dictionaries,
1352        );
1353        self.refused(admitted)?;
1354        Ok(output)
1355    }
1356
1357    fn validate_token_shapes(&mut self, input: &str) -> Result<(), CodecErrorCode> {
1358        self.transaction
1359            .validate_token_shapes(input)
1360            .map_err(|_| CodecErrorCode::RestoreFailedClosed)
1361    }
1362}
1363
1364impl PipelineRequestPseudonymizer<'_, '_, '_, '_> {
1365    fn protect_preserving_tokens(
1366        &mut self,
1367        input: &str,
1368        prefix_cache_write_mode: PrefixCacheWriteMode,
1369    ) -> Result<String, CodecErrorCode> {
1370        self.validate_token_shapes(input)?;
1371        let mut tokens = self.transaction.tokens();
1372        if tokens.is_empty() {
1373            return self.protect_unstaged_segment(input, prefix_cache_write_mode);
1374        }
1375        tokens.sort_unstable_by_key(|token| std::cmp::Reverse(token.len()));
1376
1377        let mut output = String::with_capacity(input.len());
1378        let mut cursor = 0;
1379        while cursor < input.len() {
1380            let next = tokens
1381                .iter()
1382                .filter_map(|token| {
1383                    input[cursor..]
1384                        .find(token)
1385                        .map(|offset| (cursor + offset, token.as_str()))
1386                })
1387                .min_by(|left, right| {
1388                    left.0
1389                        .cmp(&right.0)
1390                        .then_with(|| right.1.len().cmp(&left.1.len()))
1391                });
1392            let Some((start, token)) = next else {
1393                output.push_str(
1394                    &self.protect_unstaged_segment(&input[cursor..], prefix_cache_write_mode)?,
1395                );
1396                break;
1397            };
1398            if start > cursor {
1399                output.push_str(
1400                    &self
1401                        .protect_unstaged_segment(&input[cursor..start], prefix_cache_write_mode)?,
1402                );
1403            }
1404            output.push_str(token);
1405            cursor = start + token.len();
1406        }
1407        Ok(output)
1408    }
1409}
1410
1411struct PipelineResponseResidualValidator<'pipeline, 'context> {
1412    pipeline: &'pipeline Pipeline,
1413    locale_chain: &'context [LocaleTag],
1414    dictionaries: &'context DictionaryBundle,
1415    validation_session: Session,
1416}
1417
1418impl<'pipeline, 'context> PipelineResponseResidualValidator<'pipeline, 'context> {
1419    fn new(
1420        pipeline: &'pipeline Pipeline,
1421        locale_chain: &'context [LocaleTag],
1422        dictionaries: &'context DictionaryBundle,
1423    ) -> Result<Self, DirectProxyError> {
1424        let validation_session = Session::new(Scope::Ephemeral).map_err(|_| {
1425            ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::ResponseValidation)
1426        })?;
1427        Ok(Self {
1428            pipeline,
1429            locale_chain,
1430            dictionaries,
1431            validation_session,
1432        })
1433    }
1434}
1435
1436impl ResponseResidualValidator for PipelineResponseResidualValidator<'_, '_> {
1437    fn validate(
1438        &mut self,
1439        value: &str,
1440        authorized_output_ranges: &[std::ops::Range<usize>],
1441    ) -> Result<(), CodecErrorCode> {
1442        let (clean, spans, _) = self
1443            .pipeline
1444            .clean_with_safety_net_detect_context(
1445                &self.validation_session,
1446                RawDocument::Text(value.to_owned()),
1447                self.locale_chain,
1448                self.dictionaries,
1449            )
1450            .map_err(|_| CodecErrorCode::ProviderOriginPii)?;
1451        let CleanDocument::Text(clean_text) = clean else {
1452            return Err(CodecErrorCode::ProviderOriginPii);
1453        };
1454        // PRECONDITION of the range check below, not a nicety. `authorized_output_ranges` are
1455        // `value` coordinates, and the manifest is the only evidence that carries `value`
1456        // coordinates too -- so a change no manifest entry describes leaves nothing to compare.
1457        // A safety-net fallback DELETION is exactly that change: it writes
1458        // `replace_clean_span_checked(.., "", None)`, removing the bytes from the clean text
1459        // this validator discards while emitting no entry. Deciding on `spans` alone then reads
1460        // an empty manifest as "no PII found" and admits the `value` still held here verbatim.
1461        if !manifest_accounts_for_every_change(value, &clean_text, &spans) {
1462            return Err(CodecErrorCode::ProviderOriginPii);
1463        }
1464        for span in spans {
1465            if !authorized_output_ranges.iter().any(|authorized| {
1466                authorized.start <= span.raw_span.start && authorized.end >= span.raw_span.end
1467            }) {
1468                return Err(CodecErrorCode::ProviderOriginPii);
1469            }
1470        }
1471        Ok(())
1472    }
1473}
1474
1475/// True when `clean` is `raw` with each manifest entry's raw span replaced by the clean text at
1476/// that entry's clean span, and nothing else changed.
1477///
1478/// Every byte the pipeline touched is then described by an entry carrying `raw` coordinates,
1479/// which is what lets a caller decide whether that particular change was authorized. This
1480/// refuses rather than interprets whenever the entries cannot describe the pair -- reordered,
1481/// overlapping, or out-of-bounds spans included -- because the alternative is to conclude
1482/// "nothing was found" about bytes it cannot account for.
1483fn manifest_accounts_for_every_change(raw: &str, clean: &str, spans: &[EmittedTokenSpan]) -> bool {
1484    let mut ordered: Vec<&EmittedTokenSpan> = spans.iter().collect();
1485    ordered.sort_by_key(|span| span.clean_span.start);
1486    let mut raw_cursor = 0usize;
1487    let mut clean_cursor = 0usize;
1488    for span in ordered {
1489        if span.raw_span.start < raw_cursor || span.clean_span.start < clean_cursor {
1490            return false;
1491        }
1492        let (Some(raw_gap), Some(clean_gap)) = (
1493            raw.get(raw_cursor..span.raw_span.start),
1494            clean.get(clean_cursor..span.clean_span.start),
1495        ) else {
1496            return false;
1497        };
1498        if raw_gap != clean_gap {
1499            return false;
1500        }
1501        // Both ends must be real indices on their own string, or the entry does not describe
1502        // this pair at all. `get` also rejects an inverted span and a non-char-boundary end.
1503        if raw.get(span.raw_span.clone()).is_none() || clean.get(span.clean_span.clone()).is_none()
1504        {
1505            return false;
1506        }
1507        raw_cursor = span.raw_span.end;
1508        clean_cursor = span.clean_span.end;
1509    }
1510    matches!(
1511        (raw.get(raw_cursor..), clean.get(clean_cursor..)),
1512        (Some(raw_tail), Some(clean_tail)) if raw_tail == clean_tail
1513    )
1514}
1515
1516#[allow(clippy::too_many_arguments)]
1517fn prepare_and_commit_direct_request(
1518    pipeline: &Pipeline,
1519    locale_chain: &[LocaleTag],
1520    dictionaries: &DictionaryBundle,
1521    session: &Session,
1522    codec: &dyn BodyCodec,
1523    original_body: &[u8],
1524    endpoint: &Url,
1525    headers: &HeaderMap,
1526    expected_response: WireFormat,
1527    limits: CodecLimits,
1528    provider_request_projection_selected: bool,
1529    max_request_bytes: usize,
1530) -> Result<CommittedDirectRequest, DirectProxyError> {
1531    prepare_and_commit_direct_request_with_hook(
1532        pipeline,
1533        locale_chain,
1534        dictionaries,
1535        session,
1536        codec,
1537        original_body,
1538        endpoint,
1539        headers,
1540        expected_response,
1541        limits,
1542        provider_request_projection_selected,
1543        max_request_bytes,
1544        |_, _, _| Ok(()),
1545    )
1546}
1547
1548#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1549enum SessionCommitFailureKind {
1550    GenerationConflict,
1551    NonConflict,
1552}
1553
1554fn classify_session_commit_failure(error: &SessionTransactionError) -> SessionCommitFailureKind {
1555    match error {
1556        SessionTransactionError::GenerationConflict => SessionCommitFailureKind::GenerationConflict,
1557        _ => SessionCommitFailureKind::NonConflict,
1558    }
1559}
1560
1561fn map_session_commit_failure(kind: SessionCommitFailureKind) -> DirectProxyError {
1562    match kind {
1563        SessionCommitFailureKind::GenerationConflict => {
1564            ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session)
1565        }
1566        SessionCommitFailureKind::NonConflict => {
1567            ProxyErrorCode::SessionCommitFailure.error(ProxyErrorPhase::Session)
1568        }
1569    }
1570}
1571
1572#[allow(clippy::too_many_arguments)]
1573fn prepare_and_commit_direct_request_with_hook(
1574    pipeline: &Pipeline,
1575    locale_chain: &[LocaleTag],
1576    dictionaries: &DictionaryBundle,
1577    session: &Session,
1578    codec: &dyn BodyCodec,
1579    original_body: &[u8],
1580    endpoint: &Url,
1581    headers: &HeaderMap,
1582    expected_response: WireFormat,
1583    limits: CodecLimits,
1584    provider_request_projection_selected: bool,
1585    max_request_bytes: usize,
1586    mut before_commit: impl FnMut(u8, &Session, &SessionTransaction<'_>) -> Result<(), DirectProxyError>,
1587) -> Result<CommittedDirectRequest, DirectProxyError> {
1588    for attempt in 0..=1 {
1589        let mut transaction = session.begin_transaction();
1590        let proved_body = {
1591            let mut pseudonymizer = PipelineRequestPseudonymizer::new(
1592                pipeline,
1593                locale_chain,
1594                dictionaries,
1595                &mut transaction,
1596            );
1597            let protected = {
1598                let mut context =
1599                    RequestTransformContext::new(&mut pseudonymizer, WireFormat::Json, limits);
1600                if provider_request_projection_selected {
1601                    context.request_inspection_projection();
1602                }
1603                codec.protect_request(original_body, &mut context)
1604            };
1605            protected.map_err(|error| match pseudonymizer.refusal.take() {
1606                Some(refusal) if error.code() == CodecErrorCode::ProtectionFailedClosed => {
1607                    ProxyErrorCode::ProtectionRefused
1608                        .error(ProxyErrorPhase::RequestTransform)
1609                        .with_refusal(refusal)
1610                }
1611                _ => map_codec_error(error),
1612            })?
1613        };
1614        let request = DirectRequest::from_proved(
1615            endpoint.clone(),
1616            headers.clone(),
1617            proved_body,
1618            expected_response,
1619        )?;
1620        if request.body.len() > max_request_bytes {
1621            return Err(
1622                ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
1623            );
1624        }
1625        let prepared = PreparedDirectRequest::new(transaction, request);
1626        before_commit(attempt, session, &prepared.transaction)?;
1627        match prepared.commit() {
1628            Ok(committed) => return Ok(committed),
1629            Err(error) => match classify_session_commit_failure(&error) {
1630                SessionCommitFailureKind::GenerationConflict if attempt == 0 => {}
1631                kind => return Err(map_session_commit_failure(kind)),
1632            },
1633        }
1634    }
1635    Err(ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session))
1636}
1637
1638/// Direct transaction state around the only possible upstream send.
1639#[derive(Debug, Eq, PartialEq)]
1640#[allow(dead_code)]
1641pub(crate) enum DirectIoState {
1642    Prepared,
1643    CommittedNoIo,
1644    IoAttempted,
1645    ResponseComplete,
1646}
1647
1648/// Enforces one generation retry and the Prepared-to-response I/O ordering.
1649#[derive(Debug, Eq, PartialEq)]
1650#[allow(dead_code)]
1651pub(crate) struct DirectTransactionState {
1652    state: DirectIoState,
1653    generation_retries: u8,
1654}
1655
1656#[allow(dead_code)]
1657impl DirectTransactionState {
1658    #[must_use]
1659    pub(crate) const fn new() -> Self {
1660        Self {
1661            state: DirectIoState::Prepared,
1662            generation_retries: 0,
1663        }
1664    }
1665
1666    #[must_use]
1667    pub(crate) const fn state(&self) -> &DirectIoState {
1668        &self.state
1669    }
1670
1671    #[must_use]
1672    pub(crate) const fn generation_retries(&self) -> u8 {
1673        self.generation_retries
1674    }
1675
1676    #[must_use]
1677    pub(crate) const fn retains_committed_mapping(&self) -> bool {
1678        !matches!(self.state, DirectIoState::Prepared)
1679    }
1680
1681    pub(crate) fn retry_generation_conflict(&mut self) -> Result<(), DirectProxyError> {
1682        if self.state != DirectIoState::Prepared || self.generation_retries != 0 {
1683            return Err(ProxyErrorCode::SessionGenerationConflict.error(ProxyErrorPhase::Session));
1684        }
1685        self.generation_retries = 1;
1686        Ok(())
1687    }
1688
1689    pub(crate) fn mark_committed(&mut self) -> Result<(), DirectProxyError> {
1690        self.transition(DirectIoState::Prepared, DirectIoState::CommittedNoIo)
1691    }
1692
1693    pub(crate) fn mark_io_attempted(&mut self) -> Result<(), DirectProxyError> {
1694        self.transition(DirectIoState::CommittedNoIo, DirectIoState::IoAttempted)
1695    }
1696
1697    pub(crate) fn mark_response_complete(&mut self) -> Result<(), DirectProxyError> {
1698        self.transition(DirectIoState::IoAttempted, DirectIoState::ResponseComplete)
1699    }
1700
1701    fn transition(
1702        &mut self,
1703        expected: DirectIoState,
1704        next: DirectIoState,
1705    ) -> Result<(), DirectProxyError> {
1706        if self.state != expected {
1707            return Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework));
1708        }
1709        self.state = next;
1710        Ok(())
1711    }
1712
1713    fn require_committed_no_io(&self) -> Result<(), DirectProxyError> {
1714        if self.state == DirectIoState::CommittedNoIo {
1715            Ok(())
1716        } else {
1717            Err(ProxyErrorCode::InvalidStateTransition.error(ProxyErrorPhase::Framework))
1718        }
1719    }
1720}
1721
1722impl Default for DirectTransactionState {
1723    fn default() -> Self {
1724        Self::new()
1725    }
1726}
1727
1728#[derive(Clone)]
1729struct AppState {
1730    config: Arc<ProxyConfig>,
1731    pipeline: Arc<Pipeline>,
1732    client: Client,
1733    direct: Option<Arc<DirectRuntime>>,
1734    sessions: Arc<RwLock<HashMap<String, SessionEntry>>>,
1735    started_at: Instant,
1736}
1737
1738struct SessionEntry {
1739    session: Arc<LegacySession>,
1740    expires_at: Instant,
1741}
1742
1743/// A legacy-adapter conversation: its manifest plus the restore provenance the manifest does not
1744/// record.
1745struct LegacySession {
1746    session: Session,
1747    /// Tokens minted for a value captured inside a JSON string literal of a request surface, such
1748    /// as a JSON tool result or tool-call `arguments`. The manifest holds such a value in its
1749    /// JSON-escaped spelling (`\"`, `\u00fc`), which is already valid JSON string content.
1750    json_spelled_tokens: std::sync::Mutex<HashSet<String>>,
1751}
1752
1753impl LegacySession {
1754    fn new(session: Session) -> Self {
1755        Self {
1756            session,
1757            json_spelled_tokens: std::sync::Mutex::new(HashSet::new()),
1758        }
1759    }
1760
1761    /// The set only ever grows, so a panic while it was locked cannot leave it inconsistent.
1762    fn json_spelled_tokens(&self) -> std::sync::MutexGuard<'_, HashSet<String>> {
1763        self.json_spelled_tokens
1764            .lock()
1765            .unwrap_or_else(std::sync::PoisonError::into_inner)
1766    }
1767}
1768
1769#[cfg(test)]
1770static SESSION_FOR_MISS_BARRIERS: std::sync::LazyLock<
1771    std::sync::Mutex<HashMap<String, Arc<tokio::sync::Barrier>>>,
1772> = std::sync::LazyLock::new(|| std::sync::Mutex::new(HashMap::new()));
1773
1774#[derive(Debug, Clone, Serialize)]
1775#[non_exhaustive]
1776pub struct HealthSnapshot {
1777    pub uptime_secs: u64,
1778    pub bind: String,
1779    pub adapters: Vec<AdapterSnapshot>,
1780}
1781
1782#[derive(Debug, Clone, Serialize)]
1783#[non_exhaustive]
1784pub struct AdapterSnapshot {
1785    pub name: &'static str,
1786    pub upstream: String,
1787}
1788
1789/// Bind the configured address and serve until cancelled or the server fails.
1790pub async fn serve(config: ProxyConfig, pipeline: Arc<Pipeline>) -> Result<(), ProxyError> {
1791    serve_with_ready(config, pipeline, || {}).await
1792}
1793
1794/// Invoke `ready` only after configuration validation and a successful bind.
1795pub async fn serve_with_ready(
1796    config: ProxyConfig,
1797    pipeline: Arc<Pipeline>,
1798    ready: impl FnOnce(),
1799) -> Result<(), ProxyError> {
1800    let bind = config.bind;
1801    let app = build_app(config, pipeline)?;
1802    let listener = TcpListener::bind(bind)
1803        .await
1804        .map_err(|source| ProxyError::Server { source })?;
1805    ready();
1806    serve_app(listener, app).await
1807}
1808
1809/// Serve on an already bound listener, retaining ownership throughout startup.
1810///
1811/// The listener must match `config.bind` exactly, except that a configured port of
1812/// zero accepts the listener's assigned port. IP addresses (including wildcard
1813/// addresses) must match so listener ownership cannot bypass principal validation.
1814/// Health reports the actual bound address. Configuration is validated before any
1815/// request is accepted; cancellation and routing are the same as [`serve`].
1816pub async fn serve_with_listener(
1817    mut config: ProxyConfig,
1818    pipeline: Arc<Pipeline>,
1819    listener: TcpListener,
1820) -> Result<(), ProxyError> {
1821    let actual = listener
1822        .local_addr()
1823        .map_err(|source| ProxyError::Server { source })?;
1824    if config.bind.port() == 0 {
1825        config.bind.set_port(actual.port());
1826    }
1827    if config.bind != actual {
1828        return Err(ProxyError::Server {
1829            source: std::io::Error::new(
1830                std::io::ErrorKind::InvalidInput,
1831                "listener address does not match proxy configuration",
1832            ),
1833        });
1834    }
1835    let app = build_app(config, pipeline)?;
1836    serve_app(listener, app).await
1837}
1838
1839fn build_app(config: ProxyConfig, pipeline: Arc<Pipeline>) -> Result<Router, ProxyError> {
1840    let direct = DirectRuntime::from_config(&config)?.map(Arc::new);
1841    let state = AppState {
1842        config: Arc::new(config),
1843        pipeline,
1844        client: Client::new(),
1845        direct,
1846        sessions: Arc::new(RwLock::new(HashMap::new())),
1847        started_at: Instant::now(),
1848    };
1849    Ok(Router::new()
1850        .route("/_gaze_proxy/healthz", get(healthz))
1851        .fallback(proxy)
1852        .with_state(state))
1853}
1854
1855async fn serve_app(listener: TcpListener, app: Router) -> Result<(), ProxyError> {
1856    axum::serve(
1857        listener,
1858        app.into_make_service_with_connect_info::<SocketAddr>(),
1859    )
1860    .await
1861    .map_err(|source| ProxyError::Server { source })
1862}
1863
1864async fn healthz(State(state): State<AppState>) -> impl IntoResponse {
1865    axum::Json(health_snapshot(&state))
1866}
1867
1868fn health_snapshot(state: &AppState) -> HealthSnapshot {
1869    HealthSnapshot {
1870        uptime_secs: state.started_at.elapsed().as_secs(),
1871        bind: state.config.bind.to_string(),
1872        adapters: state
1873            .config
1874            .adapters
1875            .iter()
1876            .map(|adapter| AdapterSnapshot {
1877                name: adapter.name(),
1878                upstream: adapter.upstream_base().to_string(),
1879            })
1880            .collect(),
1881    }
1882}
1883
1884async fn proxy(
1885    State(state): State<AppState>,
1886    ConnectInfo(peer): ConnectInfo<SocketAddr>,
1887    request: Request<Body>,
1888) -> Response {
1889    let (parts, body) = request.into_parts();
1890    let direct_route_candidate =
1891        state.direct.is_some() && matches!(parts.uri.path(), "/v1/messages" | "/v1/messages/");
1892    let direct_ingress = if direct_route_candidate {
1893        match prepare_direct_ingress(&state, peer, &parts.method, &parts.uri, &parts.headers) {
1894            Ok(ingress) => Some(ingress),
1895            Err(error) => return direct_error_response(error),
1896        }
1897    } else {
1898        None
1899    };
1900    let is_direct = direct_ingress.is_some();
1901    let body_limit_bytes = if is_direct {
1902        state
1903            .direct
1904            .as_ref()
1905            .map_or(state.config.body_limit_bytes, |runtime| {
1906                state
1907                    .config
1908                    .body_limit_bytes
1909                    .min(runtime.client_config.max_request_bytes as u64)
1910            })
1911    } else {
1912        state.config.body_limit_bytes
1913    };
1914    let body = match collect_inbound_body(body, &parts.headers, body_limit_bytes).await {
1915        Ok(body) => body,
1916        Err(_) if is_direct => {
1917            return direct_error_response(
1918                ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation),
1919            );
1920        }
1921        Err(error) => return proxy_error_response(error),
1922    };
1923    match proxy_inner(
1924        state,
1925        parts.method,
1926        parts.uri,
1927        parts.headers,
1928        body,
1929        direct_ingress,
1930    )
1931    .await
1932    {
1933        Ok(response) => response,
1934        Err(err) => proxy_error_response(err),
1935    }
1936}
1937
1938async fn collect_inbound_body(
1939    body: Body,
1940    headers: &HeaderMap,
1941    limit_bytes: u64,
1942) -> Result<Bytes, ProxyError> {
1943    if let Some(length) = headers
1944        .get(axum::http::header::CONTENT_LENGTH)
1945        .and_then(|value| std::str::from_utf8(value.as_bytes()).ok())
1946        .and_then(|value| value.parse::<u64>().ok())
1947    {
1948        if length > limit_bytes {
1949            return Err(ProxyError::BodyTooLarge { limit_bytes });
1950        }
1951    }
1952
1953    let mut stream = body.into_data_stream();
1954    let mut collected = Vec::new();
1955    let mut accumulated = 0_u64;
1956    while let Some(chunk) = stream.next().await {
1957        let chunk = chunk.map_err(|_| ProxyError::DaemonConfig {
1958            detail: "request_body_rejected".to_string(),
1959        })?;
1960        accumulated = accumulated
1961            .checked_add(chunk.len() as u64)
1962            .ok_or(ProxyError::BodyTooLarge { limit_bytes })?;
1963        if accumulated > limit_bytes {
1964            return Err(ProxyError::BodyTooLarge { limit_bytes });
1965        }
1966        collected.extend_from_slice(&chunk);
1967    }
1968    Ok(Bytes::from(collected))
1969}
1970
1971async fn proxy_inner(
1972    state: AppState,
1973    method: Method,
1974    uri: Uri,
1975    headers: HeaderMap,
1976    body: Bytes,
1977    direct_ingress: Option<PreparedDirectIngress>,
1978) -> Result<Response, ProxyError> {
1979    if body.len() as u64 > state.config.body_limit_bytes {
1980        return Err(ProxyError::BodyTooLarge {
1981            limit_bytes: state.config.body_limit_bytes,
1982        });
1983    }
1984    let path = uri.path();
1985    let adapter = state
1986        .config
1987        .adapters
1988        .iter()
1989        .find(|adapter| adapter.matches_path(&method, path))
1990        .cloned()
1991        .ok_or_else(|| ProxyError::AdapterNotFound {
1992            path: path.to_string(),
1993            method: method.clone(),
1994        })?;
1995    let contract = adapter.contract();
1996    if let ProtocolContract::Codec(codec) = contract.protocol() {
1997        return Ok(
1998            match direct_proxy_inner(
1999                &state,
2000                direct_ingress.ok_or_else(|| ProxyError::DaemonConfig {
2001                    detail: "direct_ingress_missing".to_string(),
2002                })?,
2003                &body,
2004                codec,
2005            )
2006            .await
2007            {
2008                Ok(response) => response,
2009                Err(error) => direct_error_response(error),
2010            },
2011        );
2012    }
2013    let session = session_for(&state, &headers).await?;
2014    let mut json: Value =
2015        serde_json::from_slice(&body).map_err(|source| ProxyError::InvalidJson { source })?;
2016    // Read before redaction mutates the body: it decides how restore writes raw values into the
2017    // model's answer text.
2018    let json_output = adapter.requests_json_output(&json);
2019    // One source of truth for both passes below: widening or narrowing detection on this
2020    // path is a configuration decision, and it must land on the primary pass and the residual
2021    // re-scan identically or the two disagree about what counts as PII.
2022    let locale_chain = state.config.locale_chain();
2023    let dictionaries = state.config.dictionaries();
2024    let redacted = redact_surfaces(
2025        &state.pipeline,
2026        &session.session,
2027        adapter.request_pii_surfaces(&mut json),
2028        locale_chain,
2029        dictionaries,
2030    )?;
2031    session
2032        .json_spelled_tokens()
2033        .extend(redacted.json_spelled_tokens);
2034    // The declared coverage policy SELECTS the mechanism that establishes coverage. Reading it
2035    // here is what keeps the declaration honest: a contract cannot claim a coverage posture the
2036    // request path never acts on.
2037    match contract.coverage_policy() {
2038        // Legacy adapters carry no codec proof, so coverage is established right here, by
2039        // re-scanning the outbound body and failing closed on anything left unprotected.
2040        CoveragePolicy::LegacySurfaces => {
2041            residual_scan_request(
2042                &state.pipeline,
2043                &json,
2044                &redacted.field_paths,
2045                locale_chain,
2046                dictionaries,
2047            )?;
2048        }
2049        // A codec-proved adapter must have been routed to the codec branch above. Reaching this
2050        // point means the contract and the router disagree and no proof has been established.
2051        CoveragePolicy::CodecProved => return Err(ProxyError::UnprovenCoverage),
2052    }
2053
2054    let upstream_url = upstream_url(adapter.upstream_base(), &uri)?;
2055    let mut request = state
2056        .client
2057        .request(method.clone(), upstream_url.clone())
2058        .json(&json);
2059    request = forward_headers(request, &headers);
2060    let upstream = request
2061        .send()
2062        .await
2063        .map_err(|source| ProxyError::UpstreamUnreachable {
2064            url: upstream_url.clone(),
2065            source,
2066        })?;
2067
2068    let status = upstream.status();
2069    let upstream_headers = upstream.headers().clone();
2070    let content_type = upstream_headers
2071        .get(reqwest::header::CONTENT_TYPE)
2072        .and_then(|value| value.to_str().ok())
2073        .unwrap_or_default()
2074        .to_ascii_lowercase();
2075    let bytes = upstream
2076        .bytes()
2077        .await
2078        .map_err(|source| ProxyError::UpstreamUnreachable {
2079            url: upstream_url,
2080            source,
2081        })?;
2082
2083    let is_sse = content_type.contains("text/event-stream");
2084    let body = if is_sse {
2085        transform_sse(&adapter, &session, &bytes, json_output)?
2086    } else {
2087        let mut response_json: Value =
2088            serde_json::from_slice(&bytes).map_err(|source| ProxyError::InvalidJson { source })?;
2089        restore_surfaces(
2090            &session,
2091            adapter.response_pii_surfaces(&mut response_json),
2092            json_output,
2093        );
2094        serde_json::to_vec(&response_json).map_err(|source| ProxyError::InvalidJson { source })?
2095    };
2096
2097    let mut response = Response::builder().status(status);
2098    response = response.header(
2099        axum::http::header::CONTENT_TYPE,
2100        if is_sse {
2101            "text/event-stream"
2102        } else {
2103            "application/json"
2104        },
2105    );
2106    response
2107        .body(axum::body::Body::from(body))
2108        .map_err(|source| ProxyError::DaemonConfig {
2109            detail: source.to_string(),
2110        })
2111}
2112
2113fn prepare_direct_ingress(
2114    state: &AppState,
2115    peer: SocketAddr,
2116    method: &Method,
2117    uri: &Uri,
2118    headers: &HeaderMap,
2119) -> Result<PreparedDirectIngress, DirectProxyError> {
2120    let runtime = state
2121        .direct
2122        .as_ref()
2123        .ok_or_else(|| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Framework))?;
2124    if method != Method::POST || uri.path() != "/v1/messages" || uri.query().is_some() {
2125        return Err(ProxyErrorCode::RouteRejected.error(ProxyErrorPhase::Ingress));
2126    }
2127    let outbound_headers = validate_direct_request_headers(headers, runtime)?;
2128    let session = match runtime.session_policy {
2129        SessionPolicy::EphemeralSingleRequest => {
2130            if headers.get_all(SESSION_HEADER).iter().next().is_some() {
2131                return Err(ProxyErrorCode::UnexpectedSessionHeader
2132                    .error(ProxyErrorPhase::RequestValidation));
2133            }
2134            DirectIngressSession::Ephemeral
2135        }
2136        SessionPolicy::OpaqueHeaderContinuity(_) => {
2137            let identifier = canonical_session_identifier(headers)?;
2138            let peer_scope = if peer.ip().is_loopback() {
2139                PeerScope::Loopback
2140            } else {
2141                PeerScope::NonLoopback
2142            };
2143            let context = PrincipalContext::new(runtime.listener_scope, peer_scope, None, None)
2144                .map_err(|_| ProxyErrorCode::PrincipalRejected.error(ProxyErrorPhase::Session))?;
2145            let principal = invoke_resolver(runtime.resolver.as_ref(), context)
2146                .map_err(|_| ProxyErrorCode::PrincipalRejected.error(ProxyErrorPhase::Session))?;
2147            let registry = runtime.registry.as_ref().ok_or_else(|| {
2148                ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Session)
2149            })?;
2150            let lease = registry
2151                .acquire(principal, identifier, None)
2152                .map_err(map_registry_error)?;
2153            DirectIngressSession::Continuity(lease)
2154        }
2155    };
2156    Ok(PreparedDirectIngress {
2157        outbound_headers,
2158        session,
2159    })
2160}
2161
2162async fn direct_proxy_inner(
2163    state: &AppState,
2164    ingress: PreparedDirectIngress,
2165    body: &[u8],
2166    codec: &dyn BodyCodec,
2167) -> Result<Response, DirectProxyError> {
2168    let runtime = state
2169        .direct
2170        .as_ref()
2171        .ok_or_else(|| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Framework))?;
2172    if body.len() > runtime.client_config.max_request_bytes {
2173        return Err(
2174            ProxyErrorCode::RequestBodyLimitExceeded.error(ProxyErrorPhase::RequestValidation)
2175        );
2176    }
2177    let expected_response = parse_request_stream_format(body)?;
2178    let mut inspection = state
2179        .config
2180        .inspection
2181        .as_ref()
2182        .and_then(|producer| producer.begin_logical(runtime.inspection_endpoint_codes));
2183
2184    let mut continuity_lease = None;
2185    let mut committed = match ingress.session {
2186        DirectIngressSession::Ephemeral => {
2187            let session = Session::new(Scope::Ephemeral)
2188                .map_err(|_| ProxyErrorCode::ProxyConfiguration.error(ProxyErrorPhase::Session))?;
2189            prepare_and_commit_direct_request(
2190                &state.pipeline,
2191                state.config.locale_chain(),
2192                state.config.dictionaries(),
2193                &session,
2194                codec,
2195                body,
2196                &runtime.endpoint,
2197                &ingress.outbound_headers,
2198                expected_response,
2199                runtime.codec_limits,
2200                inspection
2201                    .as_ref()
2202                    .is_some_and(ProxyInspectionLogicalV1::provider_request_projection_selected),
2203                runtime.client_config.max_request_bytes,
2204            )?
2205        }
2206        DirectIngressSession::Continuity(lease) => {
2207            let request_guard = lease.lock_request().await;
2208            let committed = prepare_and_commit_direct_request(
2209                &state.pipeline,
2210                state.config.locale_chain(),
2211                state.config.dictionaries(),
2212                lease.session(),
2213                codec,
2214                body,
2215                &runtime.endpoint,
2216                &ingress.outbound_headers,
2217                expected_response,
2218                runtime.codec_limits,
2219                inspection
2220                    .as_ref()
2221                    .is_some_and(ProxyInspectionLogicalV1::provider_request_projection_selected),
2222                runtime.client_config.max_request_bytes,
2223            );
2224            drop(request_guard);
2225            let committed = committed?;
2226            continuity_lease = Some(lease);
2227            committed
2228        }
2229    };
2230
2231    if let Some(inspection) = inspection.as_mut() {
2232        let inspection_parsed_facts = committed.request.inspection_parsed_facts.take();
2233        inspection.emit_request_stages(
2234            body,
2235            &committed.request.body,
2236            inspection_parsed_facts,
2237            committed.request.signed_or_encrypted_surface,
2238        );
2239    }
2240
2241    if expected_response == WireFormat::Sse {
2242        let opened = runtime.client.execute_sse(committed).await?;
2243        return stage_and_replay_direct_sse(
2244            Arc::clone(&state.pipeline),
2245            state.config.locale_chain.clone(),
2246            state.config.dictionaries().clone(),
2247            opened,
2248            runtime.codec_limits,
2249            runtime.client_config.max_response_bytes,
2250            runtime.ping_interval,
2251            inspection,
2252            continuity_lease,
2253        );
2254    }
2255    let executed = runtime.client.execute(committed).await?;
2256    let response = restore_direct_response(
2257        &state.pipeline,
2258        state.config.locale_chain(),
2259        state.config.dictionaries(),
2260        codec,
2261        executed,
2262        runtime.codec_limits,
2263        inspection.as_mut(),
2264    );
2265    drop(continuity_lease);
2266    response
2267}
2268
2269#[allow(clippy::too_many_arguments)]
2270fn stage_and_replay_direct_sse(
2271    pipeline: Arc<Pipeline>,
2272    locale_chain: gaze::LocaleChain,
2273    dictionaries: DictionaryBundle,
2274    opened: CommittedDirectSse,
2275    limits: CodecLimits,
2276    max_response_bytes: usize,
2277    ping_interval: Duration,
2278    inspection: Option<ProxyInspectionLogicalV1>,
2279    continuity_lease: Option<SessionLease>,
2280) -> Result<Response, DirectProxyError> {
2281    let CommittedDirectSse {
2282        snapshot,
2283        response: upstream_response,
2284        head,
2285        lifecycle,
2286        remaining_total_timeout,
2287    } = opened;
2288    let status = head.status();
2289    let headers = head.headers().clone();
2290    let (sender, receiver) = tokio::sync::mpsc::channel::<Result<Bytes, std::io::Error>>(8);
2291    let mut response = Response::builder().status(status);
2292    for (name, value) in &headers {
2293        response = response.header(name, value);
2294    }
2295    let response = response
2296        .body(Body::from_stream(
2297            tokio_stream::wrappers::ReceiverStream::new(receiver),
2298        ))
2299        .map_err(|_| ProxyErrorCode::UpstreamProtocol.error(ProxyErrorPhase::Framework))?;
2300    tokio::spawn(async move {
2301        let _continuity_lease = continuity_lease;
2302        let mut inspection = inspection;
2303        let proof = async move {
2304            let body = collect_response_body(upstream_response, max_response_bytes).await?;
2305            validate_declared_body(WireFormat::Sse, &body, max_response_bytes)?;
2306            let mut residual = PipelineResponseResidualValidator::new(
2307                &pipeline,
2308                locale_chain.as_slice(),
2309                &dictionaries,
2310            )?;
2311            let mut context =
2312                ResponseTransformContext::new(&snapshot, &mut residual, WireFormat::Sse, limits);
2313            if inspection
2314                .as_ref()
2315                .is_some_and(ProxyInspectionLogicalV1::owner_restored_response_projection_selected)
2316            {
2317                context.request_inspection_projection();
2318            }
2319            let mut proved = crate::codecs::anthropic::AnthropicMessagesCodec
2320                .restore_response(&body, &mut context)
2321                .map_err(map_codec_error)?;
2322            if let Some(inspection) = inspection.as_mut() {
2323                let format = proved.format();
2324                let signed_or_encrypted_surface = proved.signed_or_encrypted_surface();
2325                let inspection_parsed_facts = proved.take_inspection_parsed_facts();
2326                inspection.emit_response_stages(
2327                    &body,
2328                    proved.bytes(),
2329                    format,
2330                    inspection_parsed_facts,
2331                    signed_or_encrypted_surface,
2332                );
2333            }
2334            let mut lifecycle = lifecycle;
2335            lifecycle.mark_response_complete()?;
2336            Ok::<Vec<u8>, DirectProxyError>(proved.into_bytes())
2337        };
2338        let timed_proof = tokio::time::timeout(remaining_total_timeout, proof);
2339        tokio::pin!(timed_proof);
2340        let mut ping = tokio::time::interval(ping_interval);
2341        ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2342        ping.tick().await;
2343        loop {
2344            tokio::select! {
2345                () = sender.closed() => {
2346                    break;
2347                }
2348                result = &mut timed_proof => {
2349                    let output = match result {
2350                        Ok(Ok(bytes)) => Bytes::from(bytes),
2351                        Ok(Err(_)) | Err(_) => Bytes::from_static(DIRECT_PROXY_ERROR_FRAME),
2352                    };
2353                    let _ = sender.send(Ok(output)).await;
2354                    break;
2355                }
2356                _ = ping.tick() => {
2357                    if sender
2358                        .send(Ok(Bytes::from_static(DIRECT_PROXY_PING_FRAME)))
2359                        .await
2360                        .is_err()
2361                    {
2362                        break;
2363                    }
2364                }
2365            }
2366        }
2367    });
2368    Ok(response)
2369}
2370
2371fn validate_direct_request_headers(
2372    headers: &HeaderMap,
2373    runtime: &DirectRuntime,
2374) -> Result<HeaderMap, DirectProxyError> {
2375    validate_header_space(headers, runtime.client_config)?;
2376    for singleton in [
2377        "content-type",
2378        "content-length",
2379        "transfer-encoding",
2380        "content-encoding",
2381        "connection",
2382        "x-api-key",
2383        "anthropic-version",
2384        "anthropic-beta",
2385        SESSION_HEADER,
2386    ] {
2387        reject_duplicate_request_header(headers, singleton)?;
2388    }
2389    if headers.get(axum::http::header::CONTENT_ENCODING).is_some() {
2390        return Err(
2391            ProxyErrorCode::UnsupportedContentEncoding.error(ProxyErrorPhase::RequestValidation)
2392        );
2393    }
2394    if let Some(connection) = one_request_header(headers, "connection", false)? {
2395        let connection = connection.to_str().map_err(|_| {
2396            ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation)
2397        })?;
2398        if !connection.eq_ignore_ascii_case("keep-alive")
2399            && !connection.eq_ignore_ascii_case("close")
2400        {
2401            return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation));
2402        }
2403    }
2404    if headers.get(axum::http::header::TRANSFER_ENCODING).is_some()
2405        && headers.get(axum::http::header::CONTENT_LENGTH).is_some()
2406    {
2407        return Err(ProxyErrorCode::InvalidFraming.error(ProxyErrorPhase::RequestValidation));
2408    }
2409    let content_type = one_request_header(headers, "content-type", true)?
2410        .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2411    if parse_wire_content_type(content_type)? != WireFormat::Json {
2412        return Err(ProxyErrorCode::InvalidRequestFormat.error(ProxyErrorPhase::RequestValidation));
2413    }
2414    let api_key = one_request_header(headers, "x-api-key", true)?
2415        .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2416    if !api_key
2417        .as_bytes()
2418        .iter()
2419        .all(|byte| (0x21..=0x7e).contains(byte) && *byte != b',')
2420    {
2421        return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2422    }
2423    let version = one_request_header(headers, "anthropic-version", true)?
2424        .ok_or_else(|| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2425    let version_text = std::str::from_utf8(version.as_bytes())
2426        .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2427    if !runtime
2428        .allowed_versions
2429        .iter()
2430        .any(|allowed| allowed == version_text)
2431    {
2432        return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2433    }
2434    let beta = one_request_header(headers, "anthropic-beta", false)?;
2435    if let Some(beta) = beta {
2436        let beta_text = std::str::from_utf8(beta.as_bytes()).map_err(|_| {
2437            ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation)
2438        })?;
2439        if !runtime
2440            .allowed_betas
2441            .iter()
2442            .any(|allowed| allowed == beta_text)
2443        {
2444            return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2445        }
2446    }
2447
2448    let mut outbound = HeaderMap::new();
2449    outbound.insert(
2450        axum::http::header::CONTENT_TYPE,
2451        HeaderValue::from_static("application/json"),
2452    );
2453    outbound.insert(HeaderName::from_static("x-api-key"), api_key.clone());
2454    outbound.insert(
2455        HeaderName::from_static("anthropic-version"),
2456        version.clone(),
2457    );
2458    if let Some(beta) = beta {
2459        outbound.insert(HeaderName::from_static("anthropic-beta"), beta.clone());
2460    }
2461    Ok(outbound)
2462}
2463
2464fn one_request_header<'a>(
2465    headers: &'a HeaderMap,
2466    name: &'static str,
2467    required: bool,
2468) -> Result<Option<&'a HeaderValue>, DirectProxyError> {
2469    let mut values = headers.get_all(name).iter();
2470    let first = values.next();
2471    if values.next().is_some() || first.is_some_and(|value| value.is_empty()) {
2472        return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2473    }
2474    match (first, required) {
2475        (Some(value), _) => Ok(Some(value)),
2476        (None, false) => Ok(None),
2477        (None, true) => {
2478            Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))
2479        }
2480    }
2481}
2482
2483fn reject_duplicate_request_header(
2484    headers: &HeaderMap,
2485    name: &'static str,
2486) -> Result<(), DirectProxyError> {
2487    let mut values = headers.get_all(name).iter();
2488    let _ = values.next();
2489    if values.next().is_some() {
2490        Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))
2491    } else {
2492        Ok(())
2493    }
2494}
2495
2496fn canonical_session_identifier(headers: &HeaderMap) -> Result<&str, DirectProxyError> {
2497    let value = one_request_header(headers, SESSION_HEADER, false)?
2498        .ok_or_else(|| {
2499            ProxyErrorCode::SessionIdentityRequired.error(ProxyErrorPhase::RequestValidation)
2500        })?
2501        .to_str()
2502        .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2503    let identifier = uuid::Uuid::parse_str(value)
2504        .map_err(|_| ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation))?;
2505    if identifier.get_version() != Some(uuid::Version::Random)
2506        || identifier.get_variant() != uuid::Variant::RFC4122
2507        || identifier.hyphenated().to_string() != value
2508    {
2509        return Err(ProxyErrorCode::HeaderRejected.error(ProxyErrorPhase::RequestValidation));
2510    }
2511    Ok(value)
2512}
2513
2514fn restore_direct_response(
2515    pipeline: &Pipeline,
2516    locale_chain: &[LocaleTag],
2517    dictionaries: &DictionaryBundle,
2518    codec: &dyn BodyCodec,
2519    mut executed: CommittedDirectResponse,
2520    limits: CodecLimits,
2521    inspection: Option<&mut ProxyInspectionLogicalV1>,
2522) -> Result<Response, DirectProxyError> {
2523    let (head, body) = executed.response.into_parts();
2524    let mut residual =
2525        PipelineResponseResidualValidator::new(pipeline, locale_chain, dictionaries)?;
2526    let mut context =
2527        ResponseTransformContext::new(&executed.snapshot, &mut residual, head.format(), limits);
2528    if inspection
2529        .as_deref()
2530        .is_some_and(ProxyInspectionLogicalV1::owner_restored_response_projection_selected)
2531    {
2532        context.request_inspection_projection();
2533    }
2534    let mut proved = codec
2535        .restore_response(&body, &mut context)
2536        .map_err(map_codec_error)?;
2537    if let Some(inspection) = inspection {
2538        let format = proved.format();
2539        let signed_or_encrypted_surface = proved.signed_or_encrypted_surface();
2540        let inspection_parsed_facts = proved.take_inspection_parsed_facts();
2541        inspection.emit_response_stages(
2542            &body,
2543            proved.bytes(),
2544            format,
2545            inspection_parsed_facts,
2546            signed_or_encrypted_surface,
2547        );
2548    }
2549    executed.lifecycle.mark_response_complete()?;
2550    let mut response = Response::builder().status(head.status());
2551    for (name, value) in head.headers() {
2552        response = response.header(name, value);
2553    }
2554    response
2555        .body(Body::from(proved.into_bytes()))
2556        .map_err(|_| ProxyErrorCode::UpstreamProtocol.error(ProxyErrorPhase::Framework))
2557}
2558
2559fn map_registry_error(error: RegistryError) -> DirectProxyError {
2560    let code = match error {
2561        RegistryError::InvalidSessionIdentifier => ProxyErrorCode::HeaderRejected,
2562        RegistryError::SessionExpired => ProxyErrorCode::SessionExpired,
2563        RegistryError::SessionGenerationConflict => ProxyErrorCode::SessionGenerationConflict,
2564        RegistryError::ActiveCapacityReached => ProxyErrorCode::RegistryCapacity,
2565        RegistryError::EntropyUnavailable
2566        | RegistryError::SessionCreationFailed
2567        | RegistryError::EpochExhausted
2568        | RegistryError::InternalFailure => ProxyErrorCode::ProxyConfiguration,
2569    };
2570    code.error(ProxyErrorPhase::Session)
2571}
2572
2573fn map_codec_error(error: CodecError) -> DirectProxyError {
2574    let codec_phase = error.phase();
2575    let opaque_carrier = error.opaque_carrier();
2576    let request_direction = matches!(
2577        codec_phase,
2578        CodecPhase::RequestParse | CodecPhase::RequestProtection | CodecPhase::RequestProof
2579    );
2580    let code = match error.code() {
2581        CodecErrorCode::DuplicateObjectKey => ProxyErrorCode::DuplicateObjectKey,
2582        CodecErrorCode::InternalCoverageFailure => ProxyErrorCode::InternalCoverageFailure,
2583        CodecErrorCode::ControlWouldMutate => ProxyErrorCode::ControlWouldMutate,
2584        CodecErrorCode::OpaqueMediaUninspected => ProxyErrorCode::OpaqueMediaUninspected,
2585        CodecErrorCode::SignedMutationRequired => ProxyErrorCode::SignedMutationRequired,
2586        CodecErrorCode::SignedSurfaceMalformed => ProxyErrorCode::SignedSurfaceMalformed,
2587        CodecErrorCode::InvalidSseLifecycle | CodecErrorCode::IncompleteSseStream => {
2588            ProxyErrorCode::InvalidSseLifecycle
2589        }
2590        CodecErrorCode::InvalidContentBlockIndex | CodecErrorCode::InvalidContentBlockDelta => {
2591            ProxyErrorCode::InvalidContentBlockIndex
2592        }
2593        CodecErrorCode::DepthLimitExceeded
2594        | CodecErrorCode::NodeLimitExceeded
2595        | CodecErrorCode::StringLimitExceeded
2596        | CodecErrorCode::NumberLimitExceeded
2597        | CodecErrorCode::BodyLimitExceeded
2598        | CodecErrorCode::GrowthLimitExceeded
2599        | CodecErrorCode::SseLineLimitExceeded
2600        | CodecErrorCode::SseFrameLimitExceeded
2601        | CodecErrorCode::SseEventLimitExceeded
2602        | CodecErrorCode::SseIndexLimitExceeded
2603        | CodecErrorCode::SseAccumulatorLimitExceeded => {
2604            if request_direction {
2605                ProxyErrorCode::RequestBodyLimitExceeded
2606            } else {
2607                ProxyErrorCode::ResponseBodyLimitExceeded
2608            }
2609        }
2610        CodecErrorCode::ProtectionFailedClosed
2611        | CodecErrorCode::RestoreFailedClosed
2612        | CodecErrorCode::ProviderOriginPii
2613        | CodecErrorCode::ProviderOriginNumeric
2614        | CodecErrorCode::TransformedKeyCollision
2615        | CodecErrorCode::InvalidUtf8
2616        | CodecErrorCode::InvalidJson
2617        | CodecErrorCode::InvalidWireFormat
2618        | CodecErrorCode::UnsupportedRequestSurface
2619        | CodecErrorCode::UnsupportedResponseSurface
2620        | CodecErrorCode::InvalidSseField
2621        | CodecErrorCode::InvalidSseFrame
2622        | CodecErrorCode::InvalidSseEventPair
2623        | CodecErrorCode::UnsafeFutureEvent
2624        | CodecErrorCode::ProviderErrorEvent => ProxyErrorCode::InvalidToken,
2625    };
2626    let phase = match codec_phase {
2627        CodecPhase::RequestParse => ProxyErrorPhase::RequestValidation,
2628        CodecPhase::RequestProtection | CodecPhase::RequestProof => {
2629            ProxyErrorPhase::RequestTransform
2630        }
2631        CodecPhase::ResponseParse => ProxyErrorPhase::ResponseValidation,
2632        CodecPhase::ResponseRestore | CodecPhase::ResponseProof => {
2633            ProxyErrorPhase::ResponseTransform
2634        }
2635        CodecPhase::SseDecode | CodecPhase::SseLifecycle => ProxyErrorPhase::ResponseValidation,
2636        CodecPhase::SseReplay => ProxyErrorPhase::ResponseReplay,
2637    };
2638    let direct_error = code.error(phase);
2639    if let Some(opaque_carrier) = opaque_carrier {
2640        direct_error.with_opaque_carrier(opaque_carrier)
2641    } else {
2642        direct_error
2643    }
2644}
2645
2646/// The client-facing and logged shape of a protection refusal: the typed `ProtectionError`
2647/// variant, the fallback reason, and the suspect classes. Never any request text.
2648fn refusal_json(refusal: &BoundaryRefusal) -> Value {
2649    serde_json::json!({
2650        "error": refusal.error.as_str(),
2651        "fallback_reason": refusal.fallback_reason.map(|reason| reason.as_str()),
2652        "suspect_classes": refusal
2653            .suspect_classes
2654            .iter()
2655            .map(gaze::PiiClass::to_canonical_str)
2656            .collect::<Vec<_>>(),
2657    })
2658}
2659
2660/// A refusal is an adopter-visible decision, so it goes to the proxy log (stderr) as one line.
2661fn log_refusal(refusal: &BoundaryRefusal) {
2662    eprintln!("gaze-proxy: request refused: {}", refusal_json(refusal));
2663}
2664
2665fn direct_error_response(error: DirectProxyError) -> Response {
2666    let mut error_body = serde_json::json!({
2667        "type": "api_error",
2668        "message": "proxy_validation_failed",
2669        "code": format!("{:?}", error.code()),
2670        "phase": format!("{:?}", error.phase()),
2671    });
2672    if let (Some(name), Some(reason)) = (error.header_name(), error.header_rejection_reason()) {
2673        error_body["header"] = serde_json::json!({
2674            "name": name,
2675            "reason": reason.as_str(),
2676        });
2677    }
2678    if let Some(refusal) = error.refusal() {
2679        log_refusal(refusal);
2680        error_body["refusal"] = refusal_json(refusal);
2681    }
2682    if let Some(carrier) = error.opaque_carrier() {
2683        error_body["carrier"] = serde_json::json!({
2684            "scheme": carrier.scheme().as_str(),
2685            "span": {
2686                "byte_offset": carrier.byte_offset(),
2687                "byte_length": carrier.byte_length(),
2688            },
2689            "content_block_index": carrier.content_block_index(),
2690        });
2691    }
2692    let body = serde_json::json!({
2693        "type": "error",
2694        "error": error_body,
2695    });
2696    let mut response = Response::builder()
2697        .status(error.http_status())
2698        .header(axum::http::header::CONTENT_TYPE, "application/json");
2699    if let Some(seconds) = error.retry_after_seconds() {
2700        response = response.header(axum::http::header::RETRY_AFTER, seconds.to_string());
2701    }
2702    response
2703        .body(Body::from(body.to_string()))
2704        .unwrap_or_else(|_| Response::new(Body::empty()))
2705}
2706
2707async fn session_for(
2708    state: &AppState,
2709    headers: &HeaderMap,
2710) -> Result<Arc<LegacySession>, ProxyError> {
2711    let now = Instant::now();
2712    let id = headers
2713        .get(SESSION_HEADER)
2714        .and_then(|value| value.to_str().ok())
2715        .filter(|value| !value.trim().is_empty())
2716        .map(str::to_string)
2717        .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
2718
2719    #[cfg(test)]
2720    {
2721        let barrier = SESSION_FOR_MISS_BARRIERS.lock().unwrap().get(&id).cloned();
2722        if let Some(barrier) = barrier {
2723            barrier.wait().await;
2724        }
2725    }
2726
2727    let mut sessions = state.sessions.write().await;
2728    sessions.retain(|_, entry| entry.expires_at > now);
2729    match sessions.entry(id) {
2730        std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().session.clone()),
2731        std::collections::hash_map::Entry::Vacant(entry) => {
2732            let session = Arc::new(LegacySession::new(
2733                Session::new(Scope::Conversation(entry.key().clone()))
2734                    .map_err(|source| ProxyError::Pipeline { source })?,
2735            ));
2736            entry.insert(SessionEntry {
2737                session: session.clone(),
2738                expires_at: now + state.config.session_ttl,
2739            });
2740            Ok(session)
2741        }
2742    }
2743}
2744
2745/// Numeric request positions that are caller-authored generation controls rather than data
2746/// supplied by the data owner.
2747///
2748/// [`residual_scan_request`] treats every JSON number in a legacy request as a potential PII
2749/// carrier EXCEPT numbers whose immediate key appears here. Each entry is a sampling or
2750/// generation knob: the CALLER picks the value to shape decoding, a data owner's identifier
2751/// can never arrive in it, and probing it would reject ordinary traffic — `max_tokens: 20000`
2752/// is a five-digit numeral indistinguishable from a postal code to a digit recognizer.
2753///
2754/// Schema BOUNDS (`minimum`, `maximum`, `exclusiveMinimum`, `exclusiveMaximum`, `multipleOf`)
2755/// are deliberately absent: they remain carriers and fail closed. That matches the Anthropic
2756/// codec's ruling on numeric schema literals, and a proxy that answered differently from the
2757/// codec for the same input shape would be an axis-4 determinism and traceability regression.
2758///
2759/// The list is exact and exhaustive — no wildcards, no prefix or suffix matching — so every
2760/// position exempt from probing is visible to a reviewer. Entries cover both the snake_case
2761/// (OpenAI, Anthropic) and lowerCamelCase (Gemini) spellings of the same control.
2762const NON_CARRIER_NUMERIC_KEYS: &[&str] = &[
2763    "budget_tokens",
2764    "candidateCount",
2765    "frequencyPenalty",
2766    "frequency_penalty",
2767    "logprobs",
2768    "maxOutputTokens",
2769    "max_completion_tokens",
2770    "max_tokens",
2771    "n",
2772    "presencePenalty",
2773    "presence_penalty",
2774    "seed",
2775    "temperature",
2776    "topK",
2777    "topP",
2778    "top_k",
2779    "top_logprobs",
2780    "top_p",
2781];
2782
2783/// Path segments whose subtrees keep numeric-carrier status even when a leaf key name collides
2784/// with [`NON_CARRIER_NUMERIC_KEYS`].
2785///
2786/// A tool's JSON Schema may legitimately declare a property named `seed`, `n`, or
2787/// `temperature`; a numeric literal in that position is tool-declared DATA, not a sampling
2788/// control, and must stay probed. Without this guard the non-carrier list would punch holes in
2789/// exactly the schema positions the Anthropic `const` finding proved exploitable.
2790const CARRIER_SUBTREE_KEYS: &[&str] = &[
2791    "functionDeclarations",
2792    "functions",
2793    "input_schema",
2794    "json_schema",
2795    "parameters",
2796    "properties",
2797    "response_format",
2798    "schema",
2799    "tools",
2800];
2801
2802/// Redacts every surface an adapter claims and returns the field paths it covered.
2803///
2804/// The returned set is the authorization input for [`residual_scan_request`]: a position the
2805/// adapter surfaced has passed primary policy, the configured nets' Resolve step, and
2806/// configured-net admission. Primary preserve remains allowed unless a configured net flags
2807/// it; Resolve then tokenizes it, and admission refuses anything still raw. Other positions
2808/// require the separate residual scan.
2809///
2810/// `locale_chain` is the adopter-configured chain from
2811/// [`crate::ProxyConfig::with_locale_chain`], NOT the `[LocaleTag::Global]` that
2812/// [`Pipeline::redact`] pins. Passing it is what lets locale-gated recognizers
2813/// (`postal.de`, `postal.us`, the national phone recognizers) run on proxied traffic at all.
2814/// [`residual_scan_request`] MUST be given the same slice.
2815fn redact_surfaces(
2816    pipeline: &Pipeline,
2817    session: &Session,
2818    surfaces: Vec<crate::adapter::PiiSurface<'_>>,
2819    locale_chain: &[LocaleTag],
2820    dictionaries: &DictionaryBundle,
2821) -> Result<RedactedSurfaces, ProxyError> {
2822    let mut redacted = RedactedSurfaces::default();
2823    for surface in surfaces {
2824        // `gaze clean --safety-net-fallback strict`'s Resolve step, so a net-flagged date becomes
2825        // a token here instead of a refusal at admission. Admission stays the final
2826        // gate.
2827        let text = pipeline
2828            .resolve_boundary_text(session, surface.text, locale_chain, dictionaries)
2829            .map_err(|refusal| ProxyError::Refused { refusal })?;
2830        pipeline
2831            .admit_boundary_text(session, &text, locale_chain, dictionaries)
2832            .map_err(|refusal| ProxyError::Refused { refusal })?;
2833        collect_json_spelled_tokens(surface.text, &text, &mut redacted.json_spelled_tokens);
2834        *surface.text = text;
2835        redacted.field_paths.insert(surface.field_path);
2836    }
2837    Ok(redacted)
2838}
2839
2840#[derive(Default)]
2841struct RedactedSurfaces {
2842    /// Every redacted surface's field path, which the residual re-scan skips.
2843    field_paths: HashSet<String>,
2844    /// See [`LegacySession::json_spelled_tokens`].
2845    json_spelled_tokens: HashSet<String>,
2846}
2847
2848/// Records the tokens that `clean` gained inside a JSON string literal.
2849///
2850/// When the protected text is a JSON object or array, a token inside one of its string literals
2851/// replaced a value in its JSON-escaped spelling. A token already present in `original` was not
2852/// minted here and keeps whatever spelling it has.
2853fn collect_json_spelled_tokens(original: &str, clean: &str, tokens: &mut HashSet<String>) {
2854    if !matches!(clean.trim_start().as_bytes().first(), Some(b'{' | b'[')) {
2855        return;
2856    }
2857    let Ok(document) = serde_json::from_str::<Value>(clean) else {
2858        return;
2859    };
2860    let mut stack = vec![&document];
2861    while let Some(value) = stack.pop() {
2862        let mut record = |text: &str| {
2863            for token in token_shape::find_tokens(text) {
2864                if !original.contains(token) {
2865                    tokens.insert(token.to_owned());
2866                }
2867            }
2868        };
2869        match value {
2870            Value::String(text) => record(text),
2871            Value::Array(items) => stack.extend(items),
2872            Value::Object(members) => {
2873                for (key, member) in members {
2874                    record(key);
2875                    stack.push(member);
2876                }
2877            }
2878            Value::Null | Value::Bool(_) | Value::Number(_) => {}
2879        }
2880    }
2881}
2882
2883/// Fails closed when the outbound legacy request body carries PII outside every adapter surface.
2884///
2885/// The legacy path forwards the whole original body after mutating only the surfaces an adapter
2886/// returned, so any position the adapter's allowlist does not enumerate egresses verbatim. This
2887/// re-scan states the missing invariant directly: everything we did not redact must be PII-free.
2888///
2889/// It REJECTS rather than rewriting. Silently mutating provider-owned fields would break the
2890/// provider API contract and the restore round-trip (axis 2); refusing the request preserves
2891/// axis 1 and leaves the adapter allowlist as the place to add genuine carriers.
2892///
2893/// Detection runs under the caller-supplied `locale_chain`, which MUST be the same slice
2894/// [`redact_surfaces`] was given — both read [`crate::ProxyConfig::locale_chain`]. That shared
2895/// source is what keeps the residual neither weaker NOR stronger than the primary pass: a
2896/// chain that activated a recognizer for the surface pass but not here would forward PII the
2897/// primary would have tokenized, and the reverse would reject traffic the primary accepts.
2898fn residual_scan_request(
2899    pipeline: &Pipeline,
2900    body: &Value,
2901    surfaced: &HashSet<String>,
2902    locale_chain: &[LocaleTag],
2903    dictionaries: &DictionaryBundle,
2904) -> Result<(), ProxyError> {
2905    let session =
2906        Session::new(Scope::Ephemeral).map_err(|source| ProxyError::Pipeline { source })?;
2907    let scan = RequestResidualScan {
2908        pipeline,
2909        session,
2910        surfaced,
2911        locale_chain,
2912        dictionaries,
2913    };
2914    scan.walk("", "", body, false, true)
2915}
2916
2917/// Walk state for the outbound request residual re-scan.
2918struct RequestResidualScan<'a> {
2919    pipeline: &'a Pipeline,
2920    session: Session,
2921    surfaced: &'a HashSet<String>,
2922    /// The same chain the primary pass ran under. See [`residual_scan_request`].
2923    locale_chain: &'a [LocaleTag],
2924    /// The exact bundle used by the primary pass.
2925    dictionaries: &'a DictionaryBundle,
2926}
2927
2928impl RequestResidualScan<'_> {
2929    fn probe(&self, field_path: &str, text: &str) -> Result<(), ProxyError> {
2930        let (clean, spans, report) = self
2931            .pipeline
2932            .clean_with_safety_net_detect_context(
2933                &self.session,
2934                RawDocument::Text(text.to_owned()),
2935                self.locale_chain,
2936                self.dictionaries,
2937            )
2938            .map_err(|source| ProxyError::Pipeline { source })?;
2939        // This walk forwards `text` unchanged when it returns `Ok`, so the only thing that
2940        // justifies `Ok` is the pipeline having done NOTHING to this position. An empty `spans`
2941        // does not say that: a safety-net fallback DELETION removes the PII from a clean text
2942        // that is discarded here and emits no manifest entry, so reading `spans` alone accepts
2943        // the very position a net just flagged and forwards the raw bytes to the provider.
2944        let untouched = matches!(&clean, CleanDocument::Text(cleaned) if cleaned == text);
2945        // Three conjuncts, deliberately overlapping. Under the policy this call site pins
2946        // (`SafetyNetPolicy::default()` = Resolve + Redact) a deletion always leaves BOTH a
2947        // changed text and a surviving suspect, so `untouched` and `report.suspects` are each
2948        // redundant TODAY -- a mutation probe kills neither on its own. Keep both anyway:
2949        // `untouched` states directly what `Ok` promises here (these exact bytes get forwarded),
2950        // and `report.suspects` is the only backstop left if a future mode reports without
2951        // mutating. Dropping either narrows the guard to an assumption about the current policy.
2952        if untouched && spans.is_empty() && report.suspects.is_empty() {
2953            return Ok(());
2954        }
2955        Err(ProxyError::UnsurfacedPii {
2956            field_path: field_path.to_string(),
2957        })
2958    }
2959
2960    /// Recurses the outbound body probing every scalar the adapter did not surface.
2961    ///
2962    /// `unambiguous` tracks whether the constructed path can still be compared against adapter
2963    /// field paths. A JSON key containing `.` or `[` could make two distinct nodes render the
2964    /// same path string, so once such a key is seen the subtree stops trusting path equality and
2965    /// probes unconditionally. Extra probing is always safe: a surfaced position is already clean.
2966    fn walk(
2967        &self,
2968        path: &str,
2969        leaf_key: &str,
2970        value: &Value,
2971        in_carrier_subtree: bool,
2972        unambiguous: bool,
2973    ) -> Result<(), ProxyError> {
2974        match value {
2975            Value::String(text) => {
2976                if !(unambiguous && self.surfaced.contains(path)) {
2977                    self.probe(path, text)?;
2978                }
2979            }
2980            Value::Number(number) => {
2981                if in_carrier_subtree || !NON_CARRIER_NUMERIC_KEYS.contains(&leaf_key) {
2982                    self.probe(path, &number.to_string())?;
2983                }
2984            }
2985            Value::Array(items) => {
2986                for (index, item) in items.iter().enumerate() {
2987                    self.walk(
2988                        &format!("{path}[{index}]"),
2989                        leaf_key,
2990                        item,
2991                        in_carrier_subtree,
2992                        unambiguous,
2993                    )?;
2994                }
2995            }
2996            Value::Object(map) => {
2997                for (key, child) in map {
2998                    let child_path = if path.is_empty() {
2999                        key.clone()
3000                    } else {
3001                        format!("{path}.{key}")
3002                    };
3003                    self.walk(
3004                        &child_path,
3005                        key,
3006                        child,
3007                        in_carrier_subtree || CARRIER_SUBTREE_KEYS.contains(&key.as_str()),
3008                        unambiguous && !key.contains('.') && !key.contains('['),
3009                    )?;
3010                }
3011            }
3012            Value::Null | Value::Bool(_) => {}
3013        }
3014        Ok(())
3015    }
3016}
3017
3018/// Restores every surface in its own syntax. `json_output` is the request's
3019/// [`ProviderAdapter::requests_json_output`] answer, which resolves
3020/// [`SurfaceSyntax::ModelOutput`].
3021fn restore_surfaces(
3022    session: &LegacySession,
3023    surfaces: Vec<crate::adapter::PiiSurface<'_>>,
3024    json_output: bool,
3025) {
3026    let json_spelled_tokens = session.json_spelled_tokens();
3027    for surface in surfaces {
3028        let into_json_string = match surface.syntax {
3029            SurfaceSyntax::Text => false,
3030            SurfaceSyntax::Json => true,
3031            SurfaceSyntax::ModelOutput => json_output,
3032        };
3033        *surface.text = restore_text(
3034            &session.session,
3035            &json_spelled_tokens,
3036            surface.text,
3037            into_json_string,
3038        );
3039    }
3040}
3041
3042/// Replaces every known token in `clean` with its raw value.
3043///
3044/// With `into_json_string`, `clean` is (a fragment of) a serialized JSON document. A token can
3045/// only stand inside one of its string literals, so the raw value is JSON-escaped instead of
3046/// pasted verbatim, where a `"`, `\`, or control character would break the document or change
3047/// the value the agent parses. The escape does not depend on where the literal's quotes are,
3048/// which keeps a streamed fragment correct when they arrive in a different SSE event.
3049///
3050/// A value captured inside a JSON string literal (`json_spelled_tokens`) is already escaped, so
3051/// it is written verbatim; escaping it again would add a backslash that was never in the data.
3052fn restore_text(
3053    session: &Session,
3054    json_spelled_tokens: &HashSet<String>,
3055    clean: &str,
3056    into_json_string: bool,
3057) -> String {
3058    let mut restored = String::new();
3059    let mut last = 0;
3060    for matched in token_shape::pattern().find_iter(clean) {
3061        restored.push_str(&clean[last..matched.start()]);
3062        let token = matched.as_str();
3063        match session.restore(token) {
3064            Some(raw)
3065                if into_json_string
3066                    && !(json_spelled_tokens.contains(token) && is_json_string_content(&raw)) =>
3067            {
3068                push_json_string_contents(&mut restored, &raw);
3069            }
3070            Some(raw) => restored.push_str(&raw),
3071            None => restored.push_str(token),
3072        }
3073        last = matched.end();
3074    }
3075    restored.push_str(&clean[last..]);
3076    restored
3077}
3078
3079/// Appends `raw` as the contents of a JSON string literal, without the surrounding quotes.
3080fn push_json_string_contents(out: &mut String, raw: &str) {
3081    let literal = Value::String(raw.to_owned()).to_string();
3082    out.push_str(&literal[1..literal.len() - 1]);
3083}
3084
3085/// Whether `text` can stand between the quotes of a JSON string literal as it is.
3086fn is_json_string_content(text: &str) -> bool {
3087    serde_json::from_str::<String>(&format!("\"{text}\"")).is_ok()
3088}
3089
3090fn transform_sse(
3091    adapter: &Arc<dyn ProviderAdapter>,
3092    session: &LegacySession,
3093    bytes: &[u8],
3094    json_output: bool,
3095) -> Result<Vec<u8>, ProxyError> {
3096    let text = std::str::from_utf8(bytes).map_err(|err| ProxyError::SsePartialFrame {
3097        reason: err.to_string(),
3098    })?;
3099    let mut out = String::new();
3100    for frame in text.split("\n\n") {
3101        if frame.trim().is_empty() {
3102            continue;
3103        }
3104        let mut event_name = None;
3105        let mut data_lines = Vec::new();
3106        for line in frame.lines() {
3107            if let Some(rest) = line.strip_prefix("event:") {
3108                event_name = Some(rest.trim().to_string());
3109            } else if let Some(rest) = line.strip_prefix("data:") {
3110                data_lines.push(rest.trim_start());
3111            } else {
3112                out.push_str(line);
3113                out.push('\n');
3114            }
3115        }
3116        let data = data_lines.join("\n");
3117        if data == "[DONE]" {
3118            if let Some(name) = event_name {
3119                out.push_str("event: ");
3120                out.push_str(&name);
3121                out.push('\n');
3122            }
3123            out.push_str("data: [DONE]\n\n");
3124            continue;
3125        }
3126        let mut event = SseEvent {
3127            event: event_name.clone(),
3128            data: serde_json::from_str(&data)
3129                .map_err(|source| ProxyError::InvalidJson { source })?,
3130        };
3131        restore_surfaces(
3132            session,
3133            adapter.sse_event_pii_surfaces(&mut event),
3134            json_output,
3135        );
3136        if let Some(name) = event_name {
3137            out.push_str("event: ");
3138            out.push_str(&name);
3139            out.push('\n');
3140        }
3141        out.push_str("data: ");
3142        out.push_str(
3143            &serde_json::to_string(&event.data)
3144                .map_err(|source| ProxyError::InvalidJson { source })?,
3145        );
3146        out.push_str("\n\n");
3147    }
3148    Ok(out.into_bytes())
3149}
3150
3151fn upstream_url(base: &Url, uri: &Uri) -> Result<Url, ProxyError> {
3152    let path = uri.path_and_query().map(|pq| pq.as_str()).unwrap_or("/");
3153    base.join(path).map_err(|err| ProxyError::DaemonConfig {
3154        detail: err.to_string(),
3155    })
3156}
3157
3158fn forward_headers(
3159    mut request: reqwest::RequestBuilder,
3160    headers: &HeaderMap,
3161) -> reqwest::RequestBuilder {
3162    const FORWARDED: &[&str] = &[
3163        "authorization",
3164        "x-api-key",
3165        "anthropic-version",
3166        "openai-beta",
3167        "content-type",
3168    ];
3169    for name in FORWARDED {
3170        let Ok(header_name) = HeaderName::from_bytes(name.as_bytes()) else {
3171            continue;
3172        };
3173        if let Some(value) = headers.get(&header_name) {
3174            request = request.header(name.to_string(), value.clone());
3175        }
3176    }
3177    request
3178}
3179
3180fn proxy_error_response(err: ProxyError) -> Response {
3181    let status = match err {
3182        ProxyError::AdapterNotFound { .. } => StatusCode::NOT_FOUND,
3183        ProxyError::BodyTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE,
3184        ProxyError::InvalidJson { .. } => StatusCode::BAD_REQUEST,
3185        ProxyError::UnsurfacedPii { .. } | ProxyError::Refused { .. } => {
3186            StatusCode::UNPROCESSABLE_ENTITY
3187        }
3188        ProxyError::UpstreamUnreachable { .. } => StatusCode::BAD_GATEWAY,
3189        _ => StatusCode::INTERNAL_SERVER_ERROR,
3190    };
3191    let mut body = serde_json::json!({
3192        "error": proxy_error_name(&err),
3193    });
3194    if let ProxyError::Refused { refusal } = &err {
3195        log_refusal(refusal);
3196        body["refusal"] = refusal_json(refusal);
3197    }
3198    (
3199        status,
3200        [(
3201            axum::http::header::CONTENT_TYPE,
3202            HeaderValue::from_static("application/json"),
3203        )],
3204        body.to_string(),
3205    )
3206        .into_response()
3207}
3208
3209fn proxy_error_name(err: &ProxyError) -> &'static str {
3210    match err {
3211        ProxyError::UpstreamUnreachable { .. } => "UpstreamUnreachable",
3212        ProxyError::AdapterNotFound { .. } => "AdapterNotFound",
3213        ProxyError::BodyTooLarge { .. } => "BodyTooLarge",
3214        ProxyError::InvalidJson { .. } => "InvalidJson",
3215        ProxyError::SsePartialFrame { .. } => "SsePartialFrame",
3216        ProxyError::UnsurfacedPii { .. } => "UnsurfacedPii",
3217        ProxyError::UnprovenCoverage => "UnprovenCoverage",
3218        ProxyError::Refused { .. } => "Refused",
3219        ProxyError::Pipeline { .. } => "Pipeline",
3220        ProxyError::Server { .. } => "Server",
3221        ProxyError::DaemonAlreadyRunning { .. } => "DaemonAlreadyRunning",
3222        ProxyError::DaemonNotRunning => "DaemonNotRunning",
3223        ProxyError::DaemonPidfileStale { .. } => "DaemonPidfileStale",
3224        ProxyError::DaemonIo { .. } => "DaemonIo",
3225        ProxyError::DaemonConfig { .. } => "DaemonConfig",
3226        ProxyError::DaemonExitedEarly { .. } => "DaemonExitedEarly",
3227    }
3228}
3229
3230#[cfg(test)]
3231mod tests {
3232    use super::*;
3233    use std::process::Command;
3234    use std::sync::atomic::{AtomicUsize, Ordering};
3235    use std::sync::Mutex;
3236
3237    use axum::body::to_bytes;
3238    use axum::routing::post;
3239    use gaze::{Action, ClassRule, DefaultRule, PiiClass};
3240    use gaze_recognizers::RegexDetector;
3241
3242    use crate::codec::OutputProvenance;
3243
3244    const SYNTHETIC_EMAIL: &str = "alice@example.invalid";
3245    const SYNTHETIC_PHONE: &str = "+1-555-0199";
3246
3247    struct SyntheticPrincipalResolver;
3248
3249    impl PrincipalResolver for SyntheticPrincipalResolver {
3250        fn resolve(
3251            &self,
3252            _context: PrincipalContext<'_>,
3253        ) -> Result<crate::AuthenticatedPrincipal, crate::principal::PrincipalResolveError>
3254        {
3255            crate::AuthenticatedPrincipal::try_from_canonical_bytes(b"synthetic-principal".to_vec())
3256        }
3257    }
3258
3259    #[derive(Clone, Default)]
3260    struct Capture {
3261        hits: Arc<AtomicUsize>,
3262        headers: Arc<Mutex<Option<HeaderMap>>>,
3263        body: Arc<Mutex<Vec<u8>>>,
3264    }
3265
3266    async fn capture_ok(
3267        State(capture): State<Capture>,
3268        headers: HeaderMap,
3269        body: Bytes,
3270    ) -> Response {
3271        capture.hits.fetch_add(1, Ordering::SeqCst);
3272        *capture.headers.lock().unwrap() = Some(headers);
3273        *capture.body.lock().unwrap() = body.to_vec();
3274        Response::builder()
3275            .status(StatusCode::OK)
3276            .header("content-type", "application/json")
3277            .header("set-cookie", "synthetic=drop")
3278            .header("x-request-id", "drop-me")
3279            .body(Body::from(br#"{"ok":true}"#.as_slice()))
3280            .unwrap()
3281    }
3282
3283    async fn spawn_test_server(app: Router) -> (Url, tokio::task::JoinHandle<()>) {
3284        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3285        let address = listener.local_addr().unwrap();
3286        let handle = tokio::spawn(async move {
3287            axum::serve(listener, app).await.unwrap();
3288        });
3289        (
3290            Url::parse(&format!("http://{address}/v1/messages")).unwrap(),
3291            handle,
3292        )
3293    }
3294
3295    fn proved(bytes: &[u8]) -> ProvedRequestBody {
3296        let mut provenance = OutputProvenance::default();
3297        provenance.mark_final_buffer_verified();
3298        ProvedRequestBody::new(bytes.to_vec(), WireFormat::Json, provenance)
3299    }
3300
3301    fn direct_request(
3302        url: Url,
3303        expected: WireFormat,
3304        bytes: &[u8],
3305    ) -> Result<DirectRequest, DirectProxyError> {
3306        let mut headers = HeaderMap::new();
3307        headers.insert("content-type", HeaderValue::from_static("application/json"));
3308        headers.insert("x-api-key", HeaderValue::from_static("synthetic-key"));
3309        headers.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
3310        DirectRequest::from_proved(url, headers, proved(bytes), expected)
3311    }
3312
3313    fn committed_direct_request(
3314        url: Url,
3315        expected: WireFormat,
3316        bytes: &[u8],
3317    ) -> CommittedDirectRequest {
3318        let session = Session::new(Scope::Ephemeral).unwrap();
3319        let transaction = session.begin_transaction();
3320        PreparedDirectRequest::new(transaction, direct_request(url, expected, bytes).unwrap())
3321            .commit()
3322            .unwrap()
3323    }
3324
3325    fn email_pipeline() -> Pipeline {
3326        Pipeline::builder()
3327            .detector(RegexDetector::emails().unwrap())
3328            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3329            .rule(DefaultRule::new(Action::Preserve))
3330            .build()
3331            .unwrap()
3332    }
3333
3334    #[test]
3335    fn net_admission_failure_discards_direct_staging_but_retains_legacy_live_mappings() {
3336        struct RejectNet;
3337        impl gaze::SafetyNet for RejectNet {
3338            fn id(&self) -> &str {
3339                "synthetic-reject"
3340            }
3341            fn supported_locales(&self) -> &[LocaleTag] {
3342                &[LocaleTag::Global]
3343            }
3344            fn check(
3345                &self,
3346                text: &str,
3347                _: gaze::SafetyNetContext<'_>,
3348            ) -> Result<Vec<gaze::LeakSuspect>, gaze::SafetyNetError> {
3349                if text.contains("reject-this") {
3350                    Err(gaze::SafetyNetError::Runtime {
3351                        message: "synthetic failure".into(),
3352                    })
3353                } else {
3354                    Ok(vec![])
3355                }
3356            }
3357        }
3358        let pipeline = Pipeline::builder()
3359            .detector(RegexDetector::emails().unwrap())
3360            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3361            .rule(DefaultRule::new(Action::Preserve))
3362            .register_safety_net(RejectNet)
3363            .enable_prefix_cache()
3364            .build()
3365            .unwrap();
3366        let session = Session::new(Scope::Ephemeral).unwrap();
3367        let body = format!(
3368            r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"user","content":"reject-this"}}]}}"#
3369        );
3370        let request = direct_request(
3371            Url::parse("http://127.0.0.1:1/v1/messages").unwrap(),
3372            WireFormat::Json,
3373            body.as_bytes(),
3374        )
3375        .unwrap();
3376        assert!(prepare_and_commit_direct_request(
3377            &pipeline,
3378            &[LocaleTag::Global],
3379            &DictionaryBundle::default(),
3380            &session,
3381            &crate::codecs::anthropic::AnthropicMessagesCodec,
3382            body.as_bytes(),
3383            &request.url,
3384            &request.headers,
3385            WireFormat::Json,
3386            CodecLimits::default(),
3387            false,
3388            DEFAULT_MAX_REQUEST_BYTES,
3389        )
3390        .is_err());
3391        assert!(session.tokens().is_empty());
3392        assert_eq!(session.prefix_cache_entry_count(), 0);
3393
3394        let mut first = SYNTHETIC_EMAIL.to_owned();
3395        let mut second = "reject-this".to_owned();
3396        assert!(redact_surfaces(
3397            &pipeline,
3398            &session,
3399            vec![
3400                crate::adapter::PiiSurface {
3401                    field_path: "$.first".into(),
3402                    text: &mut first,
3403                    syntax: SurfaceSyntax::Text,
3404                },
3405                crate::adapter::PiiSurface {
3406                    field_path: "$.second".into(),
3407                    text: &mut second,
3408                    syntax: SurfaceSyntax::Text,
3409                },
3410            ],
3411            &[LocaleTag::Global],
3412            &DictionaryBundle::default()
3413        )
3414        .is_err());
3415        assert_eq!(session.tokens().len(), 1);
3416        assert_eq!(
3417            session.restore_strict_text(&first).unwrap(),
3418            SYNTHETIC_EMAIL
3419        );
3420    }
3421
3422    const DIRECT_REQUEST_BODY: &[u8] = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic request"}]}"#;
3423    const DIRECT_SSE_REQUEST_BODY: &[u8] = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic request"}],"stream":true}"#;
3424    const DIRECT_JSON_RESPONSE: &[u8] = br#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[{"type":"text","text":"synthetic response"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":1,"output_tokens":1}}"#;
3425    const DIRECT_SSE_RESPONSE: &[u8] = b"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-test\",\"content\":[],\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"synthetic response\"}}\n\nevent: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":1}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
3426    const NON_STREAM_SESSION_ID: &str = "00000000-0000-4000-8000-000000000003";
3427
3428    #[derive(Clone)]
3429    struct ResponseControl {
3430        hits: Arc<AtomicUsize>,
3431        body_started: Arc<tokio::sync::Notify>,
3432        release: Arc<tokio::sync::Notify>,
3433        stream_error: bool,
3434    }
3435
3436    impl ResponseControl {
3437        fn new(stream_error: bool) -> Self {
3438            Self {
3439                hits: Arc::new(AtomicUsize::new(0)),
3440                body_started: Arc::new(tokio::sync::Notify::new()),
3441                release: Arc::new(tokio::sync::Notify::new()),
3442                stream_error,
3443            }
3444        }
3445    }
3446
3447    async fn delayed_json_response(
3448        State(control): State<ResponseControl>,
3449        _body: Bytes,
3450    ) -> Response {
3451        control.hits.fetch_add(1, Ordering::SeqCst);
3452        control.body_started.notify_one();
3453        control.release.notified().await;
3454        let body = if control.stream_error {
3455            b"{".as_slice()
3456        } else {
3457            DIRECT_JSON_RESPONSE
3458        };
3459        Response::builder()
3460            .status(StatusCode::OK)
3461            .header("content-type", "application/json")
3462            .body(Body::from(body))
3463            .unwrap()
3464    }
3465
3466    async fn controlled_sse_response(
3467        State(control): State<ResponseControl>,
3468        _body: Bytes,
3469    ) -> Response {
3470        control.hits.fetch_add(1, Ordering::SeqCst);
3471        let body_control = control.clone();
3472        let chunks = futures_util::stream::once(async move {
3473            body_control.body_started.notify_one();
3474            body_control.release.notified().await;
3475            if body_control.stream_error {
3476                Err(std::io::Error::other("synthetic stream failure"))
3477            } else {
3478                Ok(Bytes::from_static(DIRECT_SSE_RESPONSE))
3479            }
3480        });
3481        Response::builder()
3482            .status(StatusCode::OK)
3483            .header("content-type", "text/event-stream")
3484            .body(Body::from_stream(chunks))
3485            .unwrap()
3486    }
3487
3488    fn synthetic_principal() -> crate::AuthenticatedPrincipal {
3489        crate::AuthenticatedPrincipal::try_from_canonical_bytes(b"synthetic-principal".to_vec())
3490            .unwrap()
3491    }
3492
3493    fn direct_test_state(
3494        endpoint: Url,
3495        pipeline: Pipeline,
3496        session_policy: SessionPolicy,
3497        registry: Option<SessionRegistry>,
3498    ) -> AppState {
3499        let client_config = DirectClientConfig::default();
3500        let runtime = DirectRuntime {
3501            client: DirectClient::new(client_config).unwrap(),
3502            endpoint: endpoint.clone(),
3503            session_policy,
3504            registry,
3505            resolver: Arc::new(SyntheticPrincipalResolver),
3506            listener_scope: ListenerScope::Loopback,
3507            allowed_versions: Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
3508            allowed_betas: Arc::from([]),
3509            codec_limits: CodecLimits::default(),
3510            client_config,
3511            ping_interval: Duration::from_secs(10),
3512            inspection_endpoint_codes: ProxyInspectionEndpointCodesV1::from_validated_origin(
3513                &endpoint,
3514            ),
3515        };
3516        AppState {
3517            config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
3518            pipeline: Arc::new(pipeline),
3519            client: Client::new(),
3520            direct: Some(Arc::new(runtime)),
3521            sessions: Arc::new(RwLock::new(HashMap::new())),
3522            started_at: Instant::now(),
3523        }
3524    }
3525
3526    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
3527    async fn concurrent_first_requests_share_one_session() {
3528        const TASK_COUNT: usize = 8;
3529        const SESSION_ID: &str = "synthetic-concurrent-session";
3530
3531        let state = Arc::new(AppState {
3532            config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
3533            pipeline: Arc::new(Pipeline::builder().build().unwrap()),
3534            client: Client::new(),
3535            direct: None,
3536            sessions: Arc::new(RwLock::new(HashMap::new())),
3537            started_at: Instant::now(),
3538        });
3539        SESSION_FOR_MISS_BARRIERS.lock().unwrap().insert(
3540            SESSION_ID.to_string(),
3541            Arc::new(tokio::sync::Barrier::new(TASK_COUNT)),
3542        );
3543
3544        let mut tasks = Vec::with_capacity(TASK_COUNT);
3545        for _ in 0..TASK_COUNT {
3546            let state = Arc::clone(&state);
3547            tasks.push(tokio::spawn(async move {
3548                let mut headers = HeaderMap::new();
3549                headers.insert(SESSION_HEADER, HeaderValue::from_static(SESSION_ID));
3550                session_for(&state, &headers).await.unwrap()
3551            }));
3552        }
3553
3554        let mut sessions = Vec::with_capacity(TASK_COUNT);
3555        for task in tasks {
3556            sessions.push(task.await.unwrap());
3557        }
3558        SESSION_FOR_MISS_BARRIERS.lock().unwrap().remove(SESSION_ID);
3559
3560        let first = &sessions[0];
3561        assert!(
3562            sessions.iter().all(|session| Arc::ptr_eq(first, session)),
3563            "concurrent first requests returned orphaned Sessions"
3564        );
3565    }
3566
3567    fn direct_test_ingress(
3568        endpoint: Url,
3569        session: DirectIngressSession,
3570        expected: WireFormat,
3571        body: &[u8],
3572    ) -> PreparedDirectIngress {
3573        let request = direct_request(endpoint, expected, body).unwrap();
3574        PreparedDirectIngress {
3575            outbound_headers: request.headers,
3576            session,
3577        }
3578    }
3579
3580    async fn wait_for_lease_count(registry: &SessionRegistry, expected: usize) {
3581        tokio::time::timeout(Duration::from_secs(60), async {
3582            loop {
3583                if registry.stats().total_leases() == expected {
3584                    return;
3585                }
3586                tokio::task::yield_now().await;
3587            }
3588        })
3589        .await
3590        .unwrap();
3591    }
3592
3593    #[derive(Clone, Copy)]
3594    enum LimitFailureDirection {
3595        Request,
3596        Response,
3597    }
3598
3599    struct LimitFailureCodec(LimitFailureDirection);
3600
3601    impl BodyCodec for LimitFailureCodec {
3602        fn protect_request(
3603            &self,
3604            input: &[u8],
3605            _ctx: &mut RequestTransformContext<'_>,
3606        ) -> Result<ProvedRequestBody, CodecError> {
3607            match self.0 {
3608                LimitFailureDirection::Request => Err(CodecError::new(
3609                    CodecErrorCode::DepthLimitExceeded,
3610                    CodecPhase::RequestParse,
3611                )),
3612                LimitFailureDirection::Response => Ok(proved(input)),
3613            }
3614        }
3615
3616        fn restore_response(
3617            &self,
3618            _input: &[u8],
3619            _ctx: &mut ResponseTransformContext<'_>,
3620        ) -> Result<crate::codec::ProvedResponseBody, CodecError> {
3621            Err(CodecError::new(
3622                CodecErrorCode::DepthLimitExceeded,
3623                CodecPhase::ResponseParse,
3624            ))
3625        }
3626    }
3627
3628    async fn begin_delayed_non_stream(
3629        invalid_response: bool,
3630    ) -> (
3631        tokio::task::JoinHandle<Result<Response, DirectProxyError>>,
3632        ResponseControl,
3633        SessionRegistry,
3634        Instant,
3635        tokio::task::JoinHandle<()>,
3636    ) {
3637        let control = ResponseControl::new(invalid_response);
3638        let (endpoint, server) = spawn_test_server(
3639            Router::new()
3640                .route("/v1/messages", post(delayed_json_response))
3641                .with_state(control.clone()),
3642        )
3643        .await;
3644        let registry_config = crate::SessionRegistryConfig::default();
3645        let registry = SessionRegistry::with_secret(registry_config, [0x7c; 32]);
3646        let request_lease = registry
3647            .acquire(synthetic_principal(), NON_STREAM_SESSION_ID, None)
3648            .unwrap();
3649        let state = direct_test_state(
3650            endpoint.clone(),
3651            email_pipeline(),
3652            SessionPolicy::OpaqueHeaderContinuity(registry_config),
3653            Some(registry.clone()),
3654        );
3655        let ingress = direct_test_ingress(
3656            endpoint,
3657            DirectIngressSession::Continuity(request_lease),
3658            WireFormat::Json,
3659            DIRECT_REQUEST_BODY,
3660        );
3661        let request = tokio::spawn(async move {
3662            direct_proxy_inner(
3663                &state,
3664                ingress,
3665                DIRECT_REQUEST_BODY,
3666                &crate::codecs::anthropic::AnthropicMessagesCodec,
3667            )
3668            .await
3669        });
3670        control.body_started.notified().await;
3671        let expired_at = Instant::now() + Duration::from_secs(60 * 60);
3672        (request, control, registry, expired_at, server)
3673    }
3674
3675    async fn begin_controlled_sse(
3676        stream_error: bool,
3677    ) -> (
3678        Response,
3679        ResponseControl,
3680        SessionRegistry,
3681        Instant,
3682        tokio::task::JoinHandle<()>,
3683    ) {
3684        const SESSION_ID: &str = "00000000-0000-4000-8000-000000000002";
3685
3686        let control = ResponseControl::new(stream_error);
3687        let (endpoint, server) = spawn_test_server(
3688            Router::new()
3689                .route("/v1/messages", post(controlled_sse_response))
3690                .with_state(control.clone()),
3691        )
3692        .await;
3693        let registry_config = crate::SessionRegistryConfig::default();
3694        let registry = SessionRegistry::with_secret(registry_config, [0x6b; 32]);
3695        let request_lease = registry
3696            .acquire(synthetic_principal(), SESSION_ID, None)
3697            .unwrap();
3698        let state = direct_test_state(
3699            endpoint.clone(),
3700            email_pipeline(),
3701            SessionPolicy::OpaqueHeaderContinuity(registry_config),
3702            Some(registry.clone()),
3703        );
3704        let ingress = direct_test_ingress(
3705            endpoint,
3706            DirectIngressSession::Continuity(request_lease),
3707            WireFormat::Sse,
3708            DIRECT_SSE_REQUEST_BODY,
3709        );
3710
3711        let response = direct_proxy_inner(
3712            &state,
3713            ingress,
3714            DIRECT_SSE_REQUEST_BODY,
3715            &crate::codecs::anthropic::AnthropicMessagesCodec,
3716        )
3717        .await
3718        .unwrap();
3719        control.body_started.notified().await;
3720        let expired_at = Instant::now() + Duration::from_secs(60 * 60);
3721        (response, control, registry, expired_at, server)
3722    }
3723
3724    async fn assert_owned_listener_health(config: ProxyConfig, listener: TcpListener) {
3725        let actual = listener.local_addr().unwrap();
3726        let destination = SocketAddr::from(([127, 0, 0, 1], actual.port()));
3727        let server = serve_with_listener(config, Arc::new(email_pipeline()), listener);
3728        let health = async {
3729            let response = Client::new()
3730                .get(format!("http://{destination}/_gaze_proxy/healthz"))
3731                .send()
3732                .await
3733                .unwrap()
3734                .error_for_status()
3735                .unwrap()
3736                .json::<serde_json::Value>()
3737                .await
3738                .unwrap();
3739            assert_eq!(response["bind"], actual.to_string());
3740            assert_eq!(response["adapters"][0]["name"], "anthropic");
3741        };
3742        tokio::select! {
3743            result = server => panic!("owned listener server stopped: {result:?}"),
3744            result = tokio::time::timeout(Duration::from_secs(60), health) => result.unwrap(),
3745        }
3746        // Dropping the serving future releases the listener, as with `serve`.
3747    }
3748
3749    #[tokio::test]
3750    async fn reserved_port_rebind_fails_but_owned_listener_handoff_serves() {
3751        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3752        let bind = listener.local_addr().unwrap();
3753        let config = ProxyConfig::anthropic_direct(
3754            bind,
3755            AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3756        );
3757        let error = serve(config.clone(), Arc::new(email_pipeline()))
3758            .await
3759            .unwrap_err();
3760        assert!(matches!(error, ProxyError::Server { source }
3761            if source.kind() == std::io::ErrorKind::AddrInUse));
3762        assert_owned_listener_health(config, listener).await;
3763    }
3764
3765    #[tokio::test]
3766    async fn owned_listener_resolves_ephemeral_port_in_health() {
3767        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3768        let config = ProxyConfig::anthropic_direct(
3769            "127.0.0.1:0".parse().unwrap(),
3770            AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3771        );
3772        assert_owned_listener_health(config, listener).await;
3773    }
3774
3775    #[tokio::test]
3776    async fn owned_listener_rejects_wrong_ip_wildcard_and_port() {
3777        for (listen, configured) in [
3778            ("0.0.0.0:0", "127.0.0.1:0"),
3779            ("127.0.0.1:0", "0.0.0.0:0"),
3780            ("127.0.0.1:0", "127.0.0.2:0"),
3781            ("127.0.0.1:0", "127.0.0.1:1"),
3782        ] {
3783            let listener = TcpListener::bind(listen).await.unwrap();
3784            let config = ProxyConfig::anthropic_direct(
3785                configured.parse().unwrap(),
3786                AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3787            );
3788            let error = serve_with_listener(config, Arc::new(email_pipeline()), listener)
3789                .await
3790                .unwrap_err();
3791            assert!(matches!(error, ProxyError::Server { source }
3792                if source.kind() == std::io::ErrorKind::InvalidInput));
3793        }
3794    }
3795
3796    #[tokio::test]
3797    async fn owned_non_loopback_listener_requires_explicit_resolver() {
3798        let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
3799        let config = ProxyConfig::anthropic_direct(
3800            "0.0.0.0:0".parse().unwrap(),
3801            AnthropicAdapter::new(Url::parse("http://127.0.0.1:1").unwrap()),
3802        );
3803        let error = serve_with_listener(config, Arc::new(email_pipeline()), listener)
3804            .await
3805            .unwrap_err();
3806        assert!(matches!(error, ProxyError::DaemonConfig { detail }
3807            if detail == "direct_principal_required"));
3808
3809        let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
3810        let config = ProxyConfig::anthropic_direct(
3811            "0.0.0.0:0".parse().unwrap(),
3812            AnthropicAdapter::builder(Url::parse("http://127.0.0.1:1").unwrap())
3813                .principal_resolver(Arc::new(SyntheticPrincipalResolver))
3814                .build()
3815                .unwrap(),
3816        );
3817        assert_owned_listener_health(config, listener).await;
3818    }
3819
3820    #[tokio::test]
3821    async fn serve_validates_before_rebind_and_owned_listener_validates_before_serving() {
3822        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
3823        let config = ProxyConfig::anthropic_direct(
3824            listener.local_addr().unwrap(),
3825            AnthropicAdapter::new(Url::parse("http://synthetic:credential@127.0.0.1").unwrap()),
3826        );
3827        for error in [
3828            serve(config.clone(), Arc::new(email_pipeline()))
3829                .await
3830                .unwrap_err(),
3831            serve_with_listener(config, Arc::new(email_pipeline()), listener)
3832                .await
3833                .unwrap_err(),
3834        ] {
3835            assert!(matches!(error, ProxyError::DaemonConfig { detail }
3836                if detail == "direct_upstream_invalid"));
3837        }
3838    }
3839
3840    #[test]
3841    fn non_loopback_readiness_requires_an_explicit_principal_resolver() {
3842        let bind = "0.0.0.0:8787".parse().unwrap();
3843        let upstream = Url::parse("https://api.anthropic.com").unwrap();
3844        let missing = ProxyConfig::anthropic_direct(bind, AnthropicAdapter::new(upstream.clone()));
3845        assert!(DirectRuntime::from_config(&missing).is_err());
3846
3847        let configured = ProxyConfig::anthropic_direct(
3848            bind,
3849            AnthropicAdapter::builder(upstream)
3850                .principal_resolver(Arc::new(SyntheticPrincipalResolver))
3851                .build()
3852                .unwrap(),
3853        );
3854        assert!(DirectRuntime::from_config(&configured).unwrap().is_some());
3855    }
3856
3857    #[test]
3858    fn invalid_direct_upstream_fails_readiness_before_any_request() {
3859        let config = ProxyConfig::anthropic_direct(
3860            "127.0.0.1:8787".parse().unwrap(),
3861            AnthropicAdapter::new(Url::parse("https://synthetic:credential@127.0.0.1").unwrap()),
3862        );
3863        assert!(DirectRuntime::from_config(&config).is_err());
3864    }
3865
3866    #[tokio::test]
3867    async fn direct_send_requires_commit_is_single_use_and_captures_exact_proof() {
3868        let capture = Capture::default();
3869        let (url, server) = spawn_test_server(
3870            Router::new()
3871                .route("/v1/messages", post(capture_ok))
3872                .with_state(capture.clone()),
3873        )
3874        .await;
3875        let exact = br#"{"messages":[{"role":"user","content":"<Email_1>"}]}"#;
3876        let client = DirectClient::new(DirectClientConfig::default()).unwrap();
3877
3878        let session = Session::new(Scope::Ephemeral).unwrap();
3879        let prepared = PreparedDirectRequest::new(
3880            session.begin_transaction(),
3881            direct_request(url.clone(), WireFormat::Json, exact).unwrap(),
3882        );
3883        drop(prepared);
3884        assert!(session.tokens().is_empty());
3885        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
3886
3887        let mut lifecycle = DirectTransactionState::new();
3888        lifecycle.retry_generation_conflict().unwrap();
3889        assert_eq!(lifecycle.generation_retries(), 1);
3890        assert_eq!(
3891            lifecycle.retry_generation_conflict().unwrap_err().code(),
3892            ProxyErrorCode::SessionGenerationConflict
3893        );
3894        let response = client
3895            .execute(committed_direct_request(url, WireFormat::Json, exact))
3896            .await
3897            .unwrap();
3898        assert_eq!(response.lifecycle.state(), &DirectIoState::IoAttempted);
3899        assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
3900        assert_eq!(&*capture.body.lock().unwrap(), exact);
3901        assert_eq!(response.response.body(), br#"{"ok":true}"#);
3902        assert_eq!(response.response.headers().len(), 1);
3903        assert!(response.response.headers().get("set-cookie").is_none());
3904        server.abort();
3905    }
3906
3907    #[test]
3908    fn signed_shadow_probe_reuses_only_the_previously_staged_mapping() {
3909        let pipeline = Pipeline::builder()
3910            .detector(RegexDetector::emails().unwrap())
3911            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
3912            .rule(DefaultRule::new(Action::Preserve))
3913            .enable_prefix_cache()
3914            .build()
3915            .unwrap();
3916        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
3917        let session = Session::new(Scope::Ephemeral).unwrap();
3918        let expected_token = format!("<{}:Email_1>", session.session_hex());
3919        let body = format!(
3920            r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"assistant","content":[{{"type":"thinking","thinking":"{expected_token}","signature":"opaque-signature"}}]}}]}}"#
3921        );
3922        let request = direct_request(
3923            Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
3924            WireFormat::Json,
3925            body.as_bytes(),
3926        )
3927        .unwrap();
3928
3929        let committed = prepare_and_commit_direct_request_with_hook(
3930            &pipeline,
3931            &[LocaleTag::Global],
3932            &DictionaryBundle::default(),
3933            &session,
3934            &codec,
3935            body.as_bytes(),
3936            &request.url,
3937            &request.headers,
3938            WireFormat::Json,
3939            CodecLimits::default(),
3940            false,
3941            DEFAULT_MAX_REQUEST_BYTES,
3942            |attempt, live, staged| {
3943                assert_eq!(attempt, 0);
3944                assert_eq!(staged.prefix_cache_entry_count(), 0);
3945                assert_eq!(live.prefix_cache_entry_count(), 0);
3946                Ok(())
3947            },
3948        )
3949        .unwrap();
3950
3951        assert_eq!(session.tokens(), vec![expected_token.clone()]);
3952        assert_eq!(committed.snapshot.tokens(), vec![expected_token.clone()]);
3953        assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
3954        assert_eq!(session.prefix_cache_entry_count(), 0);
3955        let protected = std::str::from_utf8(&committed.request.body).unwrap();
3956        assert!(!protected.contains(SYNTHETIC_EMAIL));
3957        assert_eq!(protected.matches(&expected_token).count(), 2);
3958        assert!(protected.contains(&format!(
3959            r#""thinking":"{expected_token}","signature":"opaque-signature""#
3960        )));
3961    }
3962
3963    #[test]
3964    fn final_request_reproof_is_idempotent_for_ordinary_and_format_preserving_tokens() {
3965        let endpoint = Url::parse("https://api.anthropic.com/v1/messages").unwrap();
3966        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
3967
3968        for action in [Action::Tokenize, Action::FormatPreserve] {
3969            let pipeline = Pipeline::builder()
3970                .detector(RegexDetector::emails().unwrap())
3971                .rule(ClassRule::new(PiiClass::Email, action))
3972                .rule(DefaultRule::new(Action::Preserve))
3973                .build()
3974                .unwrap();
3975            let session = Session::new(Scope::Ephemeral).unwrap();
3976            let original = format!(
3977                r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
3978            );
3979            let headers = direct_request(endpoint.clone(), WireFormat::Json, original.as_bytes())
3980                .unwrap()
3981                .headers;
3982            let first = prepare_and_commit_direct_request(
3983                &pipeline,
3984                &[LocaleTag::Global],
3985                &DictionaryBundle::default(),
3986                &session,
3987                &codec,
3988                original.as_bytes(),
3989                &endpoint,
3990                &headers,
3991                WireFormat::Json,
3992                CodecLimits::default(),
3993                false,
3994                DEFAULT_MAX_REQUEST_BYTES,
3995            )
3996            .unwrap_or_else(|error| panic!("{action:?} first proof failed: {error:?}"));
3997            let proved_once = first.request.body.clone();
3998            assert_eq!(session.tokens().len(), 1);
3999
4000            let second = prepare_and_commit_direct_request(
4001                &pipeline,
4002                &[LocaleTag::Global],
4003                &DictionaryBundle::default(),
4004                &session,
4005                &codec,
4006                &proved_once,
4007                &endpoint,
4008                &headers,
4009                WireFormat::Json,
4010                CodecLimits::default(),
4011                false,
4012                DEFAULT_MAX_REQUEST_BYTES,
4013            )
4014            .unwrap();
4015            assert_eq!(second.request.body, proved_once);
4016            assert_eq!(session.tokens().len(), 1);
4017        }
4018    }
4019
4020    #[tokio::test]
4021    async fn late_signed_failure_discards_all_staged_state_and_does_zero_io() {
4022        let capture = Capture::default();
4023        let (url, server) = spawn_test_server(
4024            Router::new()
4025                .route("/v1/messages", post(capture_ok))
4026                .with_state(capture.clone()),
4027        )
4028        .await;
4029        let pipeline = Pipeline::builder()
4030            .detector(RegexDetector::emails().unwrap())
4031            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4032            .rule(DefaultRule::new(Action::Preserve))
4033            .enable_prefix_cache()
4034            .build()
4035            .unwrap();
4036        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4037        let session = Session::new(Scope::Ephemeral).unwrap();
4038        let body = format!(
4039            r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL}","messages":[{{"role":"assistant","content":[{{"type":"thinking","thinking":"{SYNTHETIC_EMAIL}","signature":"opaque-signature"}}]}}]}}"#
4040        );
4041        let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4042        let error = prepare_and_commit_direct_request(
4043            &pipeline,
4044            &[LocaleTag::Global],
4045            &DictionaryBundle::default(),
4046            &session,
4047            &codec,
4048            body.as_bytes(),
4049            &request.url,
4050            &request.headers,
4051            WireFormat::Json,
4052            CodecLimits::default(),
4053            false,
4054            DEFAULT_MAX_REQUEST_BYTES,
4055        )
4056        .unwrap_err();
4057
4058        assert_eq!(error.code(), ProxyErrorCode::SignedMutationRequired);
4059        assert!(session.tokens().is_empty());
4060        assert_eq!(session.prefix_cache_entry_count(), 0);
4061        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4062
4063        let valid = format!(
4064            r#"{{"model":"claude-test","max_tokens":32,"system":"{SYNTHETIC_EMAIL} reports","messages":[{{"role":"user","content":"synthetic"}}]}}"#
4065        );
4066        prepare_and_commit_direct_request(
4067            &pipeline,
4068            &[LocaleTag::Global],
4069            &DictionaryBundle::default(),
4070            &session,
4071            &codec,
4072            valid.as_bytes(),
4073            &request.url,
4074            &request.headers,
4075            WireFormat::Json,
4076            CodecLimits::default(),
4077            false,
4078            DEFAULT_MAX_REQUEST_BYTES,
4079        )
4080        .unwrap();
4081        assert_eq!(session.tokens().len(), 1);
4082        server.abort();
4083    }
4084
4085    #[tokio::test]
4086    async fn logical_domain_failure_discards_state_before_commit_and_does_zero_io() {
4087        let capture = Capture::default();
4088        let (url, server) = spawn_test_server(
4089            Router::new()
4090                .route("/v1/messages", post(capture_ok))
4091                .with_state(capture.clone()),
4092        )
4093        .await;
4094        let pipeline = Pipeline::builder()
4095            .detector(RegexDetector::emails().unwrap())
4096            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4097            .rule(DefaultRule::new(Action::Preserve))
4098            .enable_prefix_cache()
4099            .build()
4100            .unwrap();
4101        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4102        let session = Session::new(Scope::Ephemeral).unwrap();
4103        let split = 6;
4104        let (left, right) = SYNTHETIC_EMAIL.split_at(split);
4105
4106        let mut prefix_transaction = session.begin_transaction();
4107        pipeline
4108            .pseudonymize_transaction_with_detect_context(
4109                &mut prefix_transaction,
4110                RawDocument::Text(left.to_owned()),
4111                &[LocaleTag::Global],
4112                &DictionaryBundle::default(),
4113            )
4114            .unwrap();
4115        prefix_transaction.commit().unwrap();
4116        let baseline_entries = session.snapshot_entries();
4117        let baseline_tokens = session.tokens();
4118        let baseline_cache_entries = session.prefix_cache_entry_count();
4119        assert_eq!(baseline_cache_entries, 0);
4120
4121        let body = format!(
4122            r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":[{{"type":"text","text":"{left}"}},{{"type":"text","text":"{right}"}}]}}]}}"#
4123        );
4124        let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4125        let mut before_commit_calls = 0usize;
4126        let error = match prepare_and_commit_direct_request_with_hook(
4127            &pipeline,
4128            &[LocaleTag::Global],
4129            &DictionaryBundle::default(),
4130            &session,
4131            &codec,
4132            body.as_bytes(),
4133            &request.url,
4134            &request.headers,
4135            WireFormat::Json,
4136            CodecLimits::default(),
4137            false,
4138            DEFAULT_MAX_REQUEST_BYTES,
4139            |_, _, _| {
4140                before_commit_calls += 1;
4141                Ok(())
4142            },
4143        ) {
4144            Ok(_) => panic!("expected closed request proof failure"),
4145            Err(error) => error,
4146        };
4147
4148        assert_eq!(error.code(), ProxyErrorCode::ControlWouldMutate);
4149        assert_eq!(error.phase(), ProxyErrorPhase::RequestTransform);
4150        assert_eq!(error.http_status(), StatusCode::UNPROCESSABLE_ENTITY);
4151        assert_eq!(before_commit_calls, 0);
4152        assert_eq!(session.snapshot_entries(), baseline_entries);
4153        assert_eq!(session.tokens(), baseline_tokens);
4154        assert_eq!(session.prefix_cache_entry_count(), baseline_cache_entries);
4155        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4156        server.abort();
4157    }
4158
4159    #[test]
4160    fn suppression_only_request_sites_do_not_populate_the_prefix_cache() {
4161        let pipeline = Pipeline::builder()
4162            .detector(RegexDetector::emails().unwrap())
4163            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4164            .rule(DefaultRule::new(Action::Preserve))
4165            .enable_prefix_cache()
4166            .build()
4167            .unwrap();
4168        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4169        let session = Session::new(Scope::Ephemeral).unwrap();
4170        let body = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"synthetic reasoning","signature":"opaque-signature","cache_control":{"type":"ephemeral"}}]}],"tools":[{"name":"synthetic_tool","input_schema":{"type":"object","$defs":{"synthetic_definition":{"type":"string"}}}}]}"#;
4171        let request = direct_request(
4172            Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4173            WireFormat::Json,
4174            body,
4175        )
4176        .unwrap();
4177
4178        let committed = prepare_and_commit_direct_request_with_hook(
4179            &pipeline,
4180            &[LocaleTag::Global],
4181            &DictionaryBundle::default(),
4182            &session,
4183            &codec,
4184            body,
4185            &request.url,
4186            &request.headers,
4187            WireFormat::Json,
4188            CodecLimits::default(),
4189            false,
4190            DEFAULT_MAX_REQUEST_BYTES,
4191            |attempt, live, staged| {
4192                assert_eq!(attempt, 0);
4193                assert_eq!(live.prefix_cache_entry_count(), 0);
4194                assert_eq!(staged.prefix_cache_entry_count(), 0);
4195                Ok(())
4196            },
4197        )
4198        .unwrap();
4199
4200        assert!(session.tokens().is_empty());
4201        assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4202        assert_eq!(session.prefix_cache_entry_count(), 0);
4203    }
4204
4205    #[test]
4206    fn every_logical_domain_probe_family_leaves_prefix_storage_empty() {
4207        let pipeline = Pipeline::builder()
4208            .detector(RegexDetector::emails().unwrap())
4209            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4210            .rule(DefaultRule::new(Action::Preserve))
4211            .enable_prefix_cache()
4212            .build()
4213            .unwrap();
4214        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4215        let session = Session::new(Scope::Ephemeral).unwrap();
4216        let body = br#"{"model":"claude-test","max_tokens":32,"system":"system-safe","tools":[{"name":"tool-safe-name","description":"tool.safe-description","input_schema":{"type":"object","description":"schema-safe-description","properties":{"property-safe":{"type":"string"}},"required":["property-safe"]}}],"metadata":{"metadata-safe-key":"metadata-safe-value"},"messages":[{"role":"user","content":"message-safe"}]}"#;
4217        let request = direct_request(
4218            Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4219            WireFormat::Json,
4220            body,
4221        )
4222        .unwrap();
4223
4224        let committed = prepare_and_commit_direct_request_with_hook(
4225            &pipeline,
4226            &[LocaleTag::Global],
4227            &DictionaryBundle::default(),
4228            &session,
4229            &codec,
4230            body,
4231            &request.url,
4232            &request.headers,
4233            WireFormat::Json,
4234            CodecLimits::default(),
4235            false,
4236            DEFAULT_MAX_REQUEST_BYTES,
4237            |attempt, live, staged| {
4238                assert_eq!(attempt, 0);
4239                assert_eq!(live.prefix_cache_entry_count(), 0);
4240                assert_eq!(staged.prefix_cache_entry_count(), 0);
4241                Ok(())
4242            },
4243        )
4244        .unwrap();
4245
4246        assert_eq!(committed.snapshot.prefix_cache_entry_count(), 0);
4247        assert_eq!(session.prefix_cache_entry_count(), 0);
4248    }
4249
4250    struct RejectOnSecondAttemptCodec {
4251        protect_calls: AtomicUsize,
4252    }
4253
4254    impl BodyCodec for RejectOnSecondAttemptCodec {
4255        fn protect_request(
4256            &self,
4257            input: &[u8],
4258            ctx: &mut RequestTransformContext<'_>,
4259        ) -> Result<ProvedRequestBody, CodecError> {
4260            if self.protect_calls.fetch_add(1, Ordering::SeqCst) == 0 {
4261                crate::codecs::anthropic::AnthropicMessagesCodec.protect_request(input, ctx)
4262            } else {
4263                Err(CodecError::new(
4264                    CodecErrorCode::ControlWouldMutate,
4265                    CodecPhase::RequestProof,
4266                ))
4267            }
4268        }
4269
4270        fn restore_response(
4271            &self,
4272            input: &[u8],
4273            ctx: &mut ResponseTransformContext<'_>,
4274        ) -> Result<crate::codec::ProvedResponseBody, CodecError> {
4275            crate::codecs::anthropic::AnthropicMessagesCodec.restore_response(input, ctx)
4276        }
4277    }
4278
4279    #[tokio::test]
4280    async fn generation_conflict_reproofs_once_and_sends_only_the_winning_body() {
4281        let capture = Capture::default();
4282        let (url, server) = spawn_test_server(
4283            Router::new()
4284                .route("/v1/messages", post(capture_ok))
4285                .with_state(capture.clone()),
4286        )
4287        .await;
4288        let pipeline = email_pipeline();
4289        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4290        let session = Session::new(Scope::Ephemeral).unwrap();
4291        let body = format!(
4292            r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4293        );
4294        let request = direct_request(url.clone(), WireFormat::Json, body.as_bytes()).unwrap();
4295        let mut hook_calls = 0;
4296        let committed = prepare_and_commit_direct_request_with_hook(
4297            &pipeline,
4298            &[LocaleTag::Global],
4299            &DictionaryBundle::default(),
4300            &session,
4301            &codec,
4302            body.as_bytes(),
4303            &request.url,
4304            &request.headers,
4305            WireFormat::Json,
4306            CodecLimits::default(),
4307            false,
4308            DEFAULT_MAX_REQUEST_BYTES,
4309            |attempt, live, _staged| {
4310                hook_calls += 1;
4311                if attempt == 0 {
4312                    live.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4313                }
4314                Ok(())
4315            },
4316        )
4317        .unwrap();
4318        assert_eq!(hook_calls, 2);
4319        assert_eq!(session.tokens().len(), 2);
4320
4321        DirectClient::new(DirectClientConfig::default())
4322            .unwrap()
4323            .execute(committed)
4324            .await
4325            .unwrap();
4326        assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
4327        let captured = String::from_utf8(capture.body.lock().unwrap().clone()).unwrap();
4328        assert!(!captured.contains(SYNTHETIC_EMAIL));
4329        assert_eq!(captured.matches(":Email_1>").count(), 1);
4330        server.abort();
4331    }
4332
4333    #[tokio::test]
4334    async fn second_attempt_request_proof_failure_rejects_without_upstream_io() {
4335        let capture = Capture::default();
4336        let (url, server) = spawn_test_server(
4337            Router::new()
4338                .route("/v1/messages", post(capture_ok))
4339                .with_state(capture.clone()),
4340        )
4341        .await;
4342        let pipeline = Pipeline::builder()
4343            .rule(DefaultRule::new(Action::Preserve))
4344            .build()
4345            .unwrap();
4346        let codec = RejectOnSecondAttemptCodec {
4347            protect_calls: AtomicUsize::new(0),
4348        };
4349        let session = Session::new(Scope::Ephemeral).unwrap();
4350        let body = br#"{"model":"claude-test","max_tokens":32,"messages":[{"role":"user","content":"synthetic"}]}"#;
4351        let request = direct_request(url, WireFormat::Json, body).unwrap();
4352        let mut hook_calls = 0usize;
4353
4354        let error = prepare_and_commit_direct_request_with_hook(
4355            &pipeline,
4356            &[LocaleTag::Global],
4357            &DictionaryBundle::default(),
4358            &session,
4359            &codec,
4360            body,
4361            &request.url,
4362            &request.headers,
4363            WireFormat::Json,
4364            CodecLimits::default(),
4365            false,
4366            DEFAULT_MAX_REQUEST_BYTES,
4367            |attempt, live, _staged| {
4368                hook_calls += 1;
4369                assert_eq!(attempt, 0);
4370                live.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4371                Ok(())
4372            },
4373        )
4374        .unwrap_err();
4375
4376        assert_eq!(codec.protect_calls.load(Ordering::SeqCst), 2);
4377        assert_eq!(hook_calls, 1);
4378        assert_eq!(error.code(), ProxyErrorCode::ControlWouldMutate);
4379        assert_eq!(error.phase(), ProxyErrorPhase::RequestTransform);
4380        assert_eq!(error.http_status(), StatusCode::UNPROCESSABLE_ENTITY);
4381        assert_eq!(session.tokens().len(), 1);
4382        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4383        server.abort();
4384    }
4385
4386    #[tokio::test]
4387    async fn second_generation_conflict_fails_closed_without_upstream_io() {
4388        let capture = Capture::default();
4389        let (url, server) = spawn_test_server(
4390            Router::new()
4391                .route("/v1/messages", post(capture_ok))
4392                .with_state(capture.clone()),
4393        )
4394        .await;
4395        let pipeline = Pipeline::builder()
4396            .detector(RegexDetector::emails().unwrap())
4397            .rule(ClassRule::new(PiiClass::Email, Action::Tokenize))
4398            .rule(DefaultRule::new(Action::Preserve))
4399            .enable_prefix_cache()
4400            .build()
4401            .unwrap();
4402        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4403        let session = Session::new(Scope::Ephemeral).unwrap();
4404        let body = format!(
4405            r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_EMAIL}"}}]}}"#
4406        );
4407        let request = direct_request(url, WireFormat::Json, body.as_bytes()).unwrap();
4408        let error = prepare_and_commit_direct_request_with_hook(
4409            &pipeline,
4410            &[LocaleTag::Global],
4411            &DictionaryBundle::default(),
4412            &session,
4413            &codec,
4414            body.as_bytes(),
4415            &request.url,
4416            &request.headers,
4417            WireFormat::Json,
4418            CodecLimits::default(),
4419            false,
4420            DEFAULT_MAX_REQUEST_BYTES,
4421            |attempt, live, staged| {
4422                assert_eq!(staged.prefix_cache_entry_count(), 0);
4423                assert_eq!(live.prefix_cache_entry_count(), 0);
4424                live.tokenize(&PiiClass::Name, &format!("Dr. Schmidt {attempt}"))
4425                    .unwrap();
4426                Ok(())
4427            },
4428        )
4429        .unwrap_err();
4430        assert_eq!(error.code(), ProxyErrorCode::SessionGenerationConflict);
4431        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4432        assert_eq!(session.prefix_cache_entry_count(), 0);
4433        assert!(session
4434            .tokens()
4435            .iter()
4436            .all(|token| !token.contains(":Email_")));
4437        server.abort();
4438    }
4439
4440    #[tokio::test]
4441    async fn cancellation_publishes_only_after_commit_and_never_rolls_it_back() {
4442        let session = Session::new(Scope::Ephemeral).unwrap();
4443        let mut staged = session.begin_transaction();
4444        staged.tokenize(&PiiClass::Email, SYNTHETIC_EMAIL).unwrap();
4445        let prepared = PreparedDirectRequest::new(
4446            staged,
4447            direct_request(
4448                Url::parse("https://api.anthropic.com/v1/messages").unwrap(),
4449                WireFormat::Json,
4450                br#"{}"#,
4451            )
4452            .unwrap(),
4453        );
4454        drop(prepared);
4455        assert!(session.tokens().is_empty());
4456
4457        let (url, server) = spawn_test_server(Router::new().route(
4458            "/v1/messages",
4459            post(|| async {
4460                tokio::time::sleep(Duration::from_secs(1)).await;
4461                ([("content-type", "application/json")], "{}")
4462            }),
4463        ))
4464        .await;
4465        let mut staged = session.begin_transaction();
4466        staged.tokenize(&PiiClass::Email, SYNTHETIC_EMAIL).unwrap();
4467        let committed = PreparedDirectRequest::new(
4468            staged,
4469            direct_request(url, WireFormat::Json, br#"{}"#).unwrap(),
4470        )
4471        .commit()
4472        .unwrap();
4473        assert_eq!(session.tokens().len(), 1);
4474        let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4475        assert!(
4476            tokio::time::timeout(Duration::from_millis(20), client.execute(committed))
4477                .await
4478                .is_err()
4479        );
4480        assert_eq!(session.tokens().len(), 1);
4481        server.abort();
4482    }
4483
4484    #[tokio::test]
4485    async fn response_restore_uses_the_exact_committed_snapshot_after_live_mutation() {
4486        let pipeline = email_pipeline();
4487        let session = Session::new(Scope::Ephemeral).unwrap();
4488        let mut transaction = session.begin_transaction();
4489        let committed_token = transaction
4490            .tokenize(&PiiClass::Email, SYNTHETIC_EMAIL)
4491            .unwrap();
4492        let snapshot = transaction.commit().unwrap();
4493        session.tokenize(&PiiClass::Name, "Dr. Schmidt").unwrap();
4494        assert_eq!(session.tokens().len(), 2);
4495
4496        let body = format!(
4497            r#"{{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[{{"type":"text","text":"{committed_token}"}}],"stop_reason":"end_turn","stop_sequence":null,"usage":{{"input_tokens":1,"output_tokens":1}}}}"#
4498        );
4499        let mut headers = HeaderMap::new();
4500        headers.insert(
4501            axum::http::header::CONTENT_TYPE,
4502            HeaderValue::from_static("application/json"),
4503        );
4504        let executed = CommittedDirectResponse {
4505            snapshot,
4506            response: DirectResponse {
4507                head: DirectResponseHead {
4508                    status: StatusCode::OK,
4509                    format: WireFormat::Json,
4510                    headers,
4511                },
4512                body: body.into_bytes(),
4513            },
4514            lifecycle: DirectTransactionState {
4515                state: DirectIoState::IoAttempted,
4516                generation_retries: 0,
4517            },
4518        };
4519        let response = restore_direct_response(
4520            &pipeline,
4521            &[LocaleTag::Global],
4522            &DictionaryBundle::default(),
4523            &crate::codecs::anthropic::AnthropicMessagesCodec,
4524            executed,
4525            CodecLimits::default(),
4526            None,
4527        )
4528        .unwrap();
4529        let restored = to_bytes(response.into_body(), 64 * 1024).await.unwrap();
4530        assert!(std::str::from_utf8(&restored)
4531            .unwrap()
4532            .contains(SYNTHETIC_EMAIL));
4533        assert!(!std::str::from_utf8(&restored)
4534            .unwrap()
4535            .contains("Dr. Schmidt"));
4536    }
4537
4538    #[tokio::test]
4539    async fn transformed_request_overlimit_rejects_before_commit_and_zero_io() {
4540        const SESSION_ID: &str = "00000000-0000-4000-8000-000000000001";
4541
4542        let capture = Capture::default();
4543        let (url, server) = spawn_test_server(
4544            Router::new()
4545                .route("/v1/messages", post(capture_ok))
4546                .with_state(capture.clone()),
4547        )
4548        .await;
4549        let phone_class = PiiClass::Custom("phone".to_string());
4550        let pipeline = Pipeline::builder()
4551            .detector(RegexDetector::new(r"\+1-555-0199", phone_class.clone()).unwrap())
4552            .rule(ClassRule::new(phone_class, Action::Tokenize))
4553            .rule(DefaultRule::new(Action::Preserve))
4554            .enable_prefix_cache()
4555            .build()
4556            .unwrap();
4557        let codec = crate::codecs::anthropic::AnthropicMessagesCodec;
4558        let body = format!(
4559            r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_PHONE}"}}]}}"#
4560        );
4561        let request_limit = body.len() + 1;
4562        let probe_session = Session::new(Scope::Ephemeral).unwrap();
4563        let probe_request = direct_request(url.clone(), WireFormat::Json, body.as_bytes()).unwrap();
4564        let transformed = prepare_and_commit_direct_request(
4565            &pipeline,
4566            &[LocaleTag::Global],
4567            &DictionaryBundle::default(),
4568            &probe_session,
4569            &codec,
4570            body.as_bytes(),
4571            &probe_request.url,
4572            &probe_request.headers,
4573            WireFormat::Json,
4574            CodecLimits::default(),
4575            false,
4576            DEFAULT_MAX_REQUEST_BYTES,
4577        )
4578        .unwrap();
4579        assert!(body.len() < request_limit);
4580        assert!(transformed.request.body.len() > request_limit);
4581
4582        let registry_config = crate::SessionRegistryConfig::default();
4583        let registry = SessionRegistry::with_secret(registry_config, [0x5a; 32]);
4584        let observed_lease = registry
4585            .acquire(
4586                crate::AuthenticatedPrincipal::try_from_canonical_bytes(
4587                    b"synthetic-principal".to_vec(),
4588                )
4589                .unwrap(),
4590                SESSION_ID,
4591                None,
4592            )
4593            .unwrap();
4594        let request_lease = registry
4595            .acquire(
4596                crate::AuthenticatedPrincipal::try_from_canonical_bytes(
4597                    b"synthetic-principal".to_vec(),
4598                )
4599                .unwrap(),
4600                SESSION_ID,
4601                None,
4602            )
4603            .unwrap();
4604        let client_config = DirectClientConfig::default()
4605            .try_with_max_request_bytes(request_limit)
4606            .unwrap();
4607        let runtime = DirectRuntime {
4608            client: DirectClient::new(client_config).unwrap(),
4609            endpoint: url.clone(),
4610            session_policy: SessionPolicy::OpaqueHeaderContinuity(registry_config),
4611            registry: Some(registry),
4612            resolver: Arc::new(SyntheticPrincipalResolver),
4613            listener_scope: ListenerScope::Loopback,
4614            allowed_versions: Arc::from([DEFAULT_ANTHROPIC_VERSION.to_string()]),
4615            allowed_betas: Arc::from([]),
4616            codec_limits: CodecLimits::default(),
4617            client_config,
4618            ping_interval: Duration::from_secs(10),
4619            inspection_endpoint_codes: ProxyInspectionEndpointCodesV1::from_validated_origin(&url),
4620        };
4621        let state = AppState {
4622            config: Arc::new(ProxyConfig::new("127.0.0.1:0".parse().unwrap(), Vec::new())),
4623            pipeline: Arc::new(pipeline),
4624            client: Client::new(),
4625            direct: Some(Arc::new(runtime)),
4626            sessions: Arc::new(RwLock::new(HashMap::new())),
4627            started_at: Instant::now(),
4628        };
4629        let ingress = PreparedDirectIngress {
4630            outbound_headers: probe_request.headers.clone(),
4631            session: DirectIngressSession::Continuity(request_lease),
4632        };
4633
4634        let error = direct_proxy_inner(&state, ingress, body.as_bytes(), &codec)
4635            .await
4636            .unwrap_err();
4637
4638        assert_eq!(error.code(), ProxyErrorCode::RequestBodyLimitExceeded);
4639        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4640        assert!(observed_lease.session().tokens().is_empty());
4641        assert!(observed_lease.session().snapshot_entries().is_empty());
4642
4643        let extended_body = format!(
4644            r#"{{"model":"claude-test","max_tokens":32,"messages":[{{"role":"user","content":"{SYNTHETIC_PHONE} reports"}}]}}"#
4645        );
4646        let recovered = prepare_and_commit_direct_request(
4647            &state.pipeline,
4648            state.config.locale_chain(),
4649            state.config.dictionaries(),
4650            observed_lease.session(),
4651            &codec,
4652            extended_body.as_bytes(),
4653            &url,
4654            &probe_request.headers,
4655            WireFormat::Json,
4656            CodecLimits::default(),
4657            false,
4658            DEFAULT_MAX_REQUEST_BYTES,
4659        )
4660        .unwrap();
4661        let first_token = format!(
4662            "<{}:Custom:phone_1>",
4663            observed_lease.session().session_hex()
4664        );
4665        assert_eq!(observed_lease.session().tokens(), vec![first_token.clone()]);
4666        assert!(std::str::from_utf8(&recovered.request.body)
4667            .unwrap()
4668            .contains(&first_token));
4669        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4670        server.abort();
4671    }
4672
4673    #[tokio::test]
4674    async fn overlimit_proved_request_and_duplicate_credentials_do_zero_io() {
4675        let capture = Capture::default();
4676        let (url, server) = spawn_test_server(
4677            Router::new()
4678                .route("/v1/messages", post(capture_ok))
4679                .with_state(capture.clone()),
4680        )
4681        .await;
4682        let config = DirectClientConfig::default()
4683            .try_with_max_request_bytes(4)
4684            .unwrap();
4685        let client = DirectClient::new(config).unwrap();
4686        let error = client
4687            .execute(committed_direct_request(
4688                url.clone(),
4689                WireFormat::Json,
4690                br#"{"over":true}"#,
4691            ))
4692            .await
4693            .unwrap_err();
4694        assert_eq!(error.code(), ProxyErrorCode::RequestBodyLimitExceeded);
4695        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4696
4697        for singleton in ["x-api-key", "anthropic-version", "anthropic-beta"] {
4698            let mut duplicate = HeaderMap::new();
4699            duplicate.insert("content-type", HeaderValue::from_static("application/json"));
4700            duplicate.append(singleton, HeaderValue::from_static("synthetic-value-1"));
4701            duplicate.append(singleton, HeaderValue::from_static("synthetic-value-2"));
4702            let error = DirectRequest::from_proved(
4703                url.clone(),
4704                duplicate,
4705                proved(br#"{}"#),
4706                WireFormat::Json,
4707            )
4708            .err()
4709            .unwrap();
4710            assert_eq!(error.code(), ProxyErrorCode::HeaderRejected);
4711        }
4712
4713        let mut unverified_headers = HeaderMap::new();
4714        unverified_headers.insert("content-type", HeaderValue::from_static("application/json"));
4715        let unverified = ProvedRequestBody::new(
4716            br#"{}"#.to_vec(),
4717            WireFormat::Json,
4718            OutputProvenance::default(),
4719        );
4720        let error =
4721            DirectRequest::from_proved(url, unverified_headers, unverified, WireFormat::Json)
4722                .err()
4723                .unwrap();
4724        assert_eq!(error.code(), ProxyErrorCode::InvalidProvenance);
4725        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
4726        server.abort();
4727    }
4728
4729    #[tokio::test]
4730    async fn header_rejections_name_only_the_header_and_disambiguate_the_cause() {
4731        const SECRET_ONE: &str = "synthetic-credential-one";
4732        const SECRET_TWO: &str = "synthetic-credential-two";
4733
4734        let mut not_allowlisted = HeaderMap::new();
4735        not_allowlisted.insert("x-synthetic-extra", HeaderValue::from_static(SECRET_ONE));
4736
4737        let mut empty = HeaderMap::new();
4738        empty.insert("anthropic-beta", HeaderValue::from_static(""));
4739
4740        let mut duplicate = HeaderMap::new();
4741        duplicate.append("x-api-key", HeaderValue::from_static(SECRET_ONE));
4742        duplicate.append("x-api-key", HeaderValue::from_static(SECRET_TWO));
4743
4744        let cases = [
4745            (
4746                not_allowlisted,
4747                "x-synthetic-extra",
4748                HeaderRejectionReason::NameNotAllowlisted,
4749            ),
4750            (empty, "anthropic-beta", HeaderRejectionReason::EmptyValue),
4751            (
4752                duplicate,
4753                "x-api-key",
4754                HeaderRejectionReason::DuplicateSingleton,
4755            ),
4756        ];
4757
4758        for (headers, expected_name, expected_reason) in cases {
4759            let error = validate_outbound_headers(&headers).unwrap_err();
4760            assert_eq!(error.code(), ProxyErrorCode::HeaderRejected);
4761            assert_eq!(error.phase(), ProxyErrorPhase::RequestValidation);
4762            assert_eq!(error.header_name().as_deref(), Some(expected_name));
4763            assert_eq!(error.header_rejection_reason(), Some(expected_reason));
4764
4765            let debug = format!("{error:?}");
4766            assert!(debug.contains(expected_name));
4767            assert!(!debug.contains(SECRET_ONE));
4768            assert!(!debug.contains(SECRET_TWO));
4769
4770            let response = direct_error_response(error);
4771            let bytes = to_bytes(response.into_body(), 4 * 1024).await.unwrap();
4772            let body: Value = serde_json::from_slice(&bytes).unwrap();
4773            assert_eq!(body["error"]["phase"], "RequestValidation");
4774            assert_eq!(body["error"]["header"]["name"], expected_name);
4775            assert_eq!(body["error"]["header"]["reason"], expected_reason.as_str());
4776            let rendered = String::from_utf8(bytes.to_vec()).unwrap();
4777            assert!(!rendered.contains(SECRET_ONE));
4778            assert!(!rendered.contains(SECRET_TWO));
4779        }
4780
4781        let maximum_name = format!("x-{}", "a".repeat(126));
4782        let mut maximum_length = HeaderMap::new();
4783        maximum_length.insert(
4784            HeaderName::from_bytes(maximum_name.as_bytes()).unwrap(),
4785            HeaderValue::from_static(SECRET_ONE),
4786        );
4787        let error = validate_outbound_headers(&maximum_length).unwrap_err();
4788        assert_eq!(error.header_name().as_deref(), Some(maximum_name.as_str()));
4789        assert_eq!(
4790            error.header_rejection_reason(),
4791            Some(HeaderRejectionReason::NameNotAllowlisted)
4792        );
4793    }
4794
4795    #[test]
4796    fn actionable_header_reporting_preserves_the_rejected_request_set() {
4797        let mut valid = HeaderMap::new();
4798        valid.insert("content-type", HeaderValue::from_static("application/json"));
4799        valid.insert("x-api-key", HeaderValue::from_static("synthetic-key"));
4800        valid.insert("anthropic-version", HeaderValue::from_static("2023-06-01"));
4801
4802        let mut not_allowlisted = valid.clone();
4803        not_allowlisted.insert(
4804            "x-synthetic-extra",
4805            HeaderValue::from_static("synthetic-value"),
4806        );
4807
4808        let mut empty = valid.clone();
4809        empty.insert("anthropic-version", HeaderValue::from_static(""));
4810
4811        let mut duplicate = valid.clone();
4812        duplicate.append("x-api-key", HeaderValue::from_static("synthetic-key-two"));
4813
4814        let rejected = [valid, not_allowlisted, empty, duplicate]
4815            .iter()
4816            .map(|headers| validate_outbound_headers(headers).is_err())
4817            .collect::<Vec<_>>();
4818        assert_eq!(rejected, vec![false, true, true, true]);
4819    }
4820
4821    #[tokio::test]
4822    async fn redirects_retries_and_response_overflow_remain_disabled() {
4823        let target_hits = Arc::new(AtomicUsize::new(0));
4824        let target_hits_for_handler = Arc::clone(&target_hits);
4825        let (target_url, target) = spawn_test_server(Router::new().route(
4826            "/v1/messages",
4827            post(move || {
4828                let hits = Arc::clone(&target_hits_for_handler);
4829                async move {
4830                    hits.fetch_add(1, Ordering::SeqCst);
4831                    StatusCode::OK
4832                }
4833            }),
4834        ))
4835        .await;
4836        let origin_hits = Arc::new(AtomicUsize::new(0));
4837        let origin_hits_for_handler = Arc::clone(&origin_hits);
4838        let location = target_url.to_string();
4839        let (origin_url, origin) = spawn_test_server(Router::new().route(
4840            "/v1/messages",
4841            post(move || {
4842                let hits = Arc::clone(&origin_hits_for_handler);
4843                let location = location.clone();
4844                async move {
4845                    hits.fetch_add(1, Ordering::SeqCst);
4846                    Response::builder()
4847                        .status(StatusCode::TEMPORARY_REDIRECT)
4848                        .header("location", location)
4849                        .body(Body::from("synthetic redirect body"))
4850                        .unwrap()
4851                }
4852            }),
4853        ))
4854        .await;
4855        let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4856        let error = client
4857            .execute(committed_direct_request(
4858                origin_url,
4859                WireFormat::Json,
4860                br#"{}"#,
4861            ))
4862            .await
4863            .unwrap_err();
4864        assert_eq!(error.code(), ProxyErrorCode::UpstreamRedirect);
4865        assert_eq!(origin_hits.load(Ordering::SeqCst), 1);
4866        assert_eq!(target_hits.load(Ordering::SeqCst), 0);
4867        origin.abort();
4868        target.abort();
4869
4870        let attempts = Arc::new(AtomicUsize::new(0));
4871        let attempts_for_handler = Arc::clone(&attempts);
4872        let (url, server) = spawn_test_server(Router::new().route(
4873            "/v1/messages",
4874            post(move || {
4875                let attempts = Arc::clone(&attempts_for_handler);
4876                async move {
4877                    attempts.fetch_add(1, Ordering::SeqCst);
4878                    StatusCode::SERVICE_UNAVAILABLE
4879                }
4880            }),
4881        ))
4882        .await;
4883        let error = client
4884            .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
4885            .await
4886            .unwrap_err();
4887        assert_eq!(error.code(), ProxyErrorCode::UpstreamUnavailable);
4888        assert_eq!(attempts.load(Ordering::SeqCst), 1);
4889        server.abort();
4890
4891        let chunks = futures_util::stream::iter([
4892            Ok::<_, std::convert::Infallible>(Bytes::from_static(b"12345")),
4893            Ok(Bytes::from_static(b"67890")),
4894        ]);
4895        let body = Arc::new(Mutex::new(Some(Body::from_stream(chunks))));
4896        let body_for_handler = Arc::clone(&body);
4897        let (url, server) = spawn_test_server(Router::new().route(
4898            "/v1/messages",
4899            post(move || {
4900                let body = body_for_handler.lock().unwrap().take().unwrap();
4901                async move {
4902                    Response::builder()
4903                        .status(StatusCode::OK)
4904                        .header("content-type", "application/json")
4905                        .body(body)
4906                        .unwrap()
4907                }
4908            }),
4909        ))
4910        .await;
4911        let config = DirectClientConfig::default()
4912            .try_with_max_response_bytes(8)
4913            .unwrap();
4914        let client = DirectClient::new(config).unwrap();
4915        let error = client
4916            .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
4917            .await
4918            .unwrap_err();
4919        assert_eq!(error.code(), ProxyErrorCode::ResponseBodyLimitExceeded);
4920        server.abort();
4921    }
4922
4923    #[tokio::test]
4924    async fn timeout_and_body_stream_error_retain_io_attempted() {
4925        let (slow_url, slow_server) = spawn_test_server(Router::new().route(
4926            "/v1/messages",
4927            post(|| async {
4928                tokio::time::sleep(Duration::from_millis(100)).await;
4929                ([("content-type", "application/json")], "{}")
4930            }),
4931        ))
4932        .await;
4933        let config = DirectClientConfig::default()
4934            .try_with_timeouts(
4935                Duration::from_millis(20),
4936                Duration::from_millis(80),
4937                Duration::from_millis(30),
4938            )
4939            .unwrap();
4940        let client = DirectClient::new(config).unwrap();
4941        let error = client
4942            .execute(committed_direct_request(
4943                slow_url,
4944                WireFormat::Json,
4945                br#"{}"#,
4946            ))
4947            .await
4948            .unwrap_err();
4949        assert_eq!(error.code(), ProxyErrorCode::TotalTimeout);
4950        slow_server.abort();
4951
4952        let chunks =
4953            futures_util::stream::once(async { Ok::<_, std::io::Error>(Bytes::from_static(b"{")) })
4954                .chain(futures_util::stream::once(async {
4955                    tokio::time::sleep(Duration::from_millis(20)).await;
4956                    Err(std::io::Error::other("synthetic stream failure"))
4957                }));
4958        let body = Body::from_stream(chunks);
4959        let body = Arc::new(Mutex::new(Some(body)));
4960        let body_for_handler = Arc::clone(&body);
4961        let (url, server) = spawn_test_server(Router::new().route(
4962            "/v1/messages",
4963            post(move || {
4964                let body = body_for_handler.lock().unwrap().take().unwrap();
4965                async move {
4966                    Response::builder()
4967                        .status(StatusCode::OK)
4968                        .header("content-type", "application/json")
4969                        .body(body)
4970                        .unwrap()
4971                }
4972            }),
4973        ))
4974        .await;
4975        let client = DirectClient::new(DirectClientConfig::default()).unwrap();
4976        let error = client
4977            .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
4978            .await
4979            .unwrap_err();
4980        assert_eq!(error.phase(), ProxyErrorPhase::UpstreamBody);
4981        server.abort();
4982    }
4983
4984    #[tokio::test]
4985    async fn codec_limit_direction_pins_code_phase_status_body_and_io() {
4986        let capture = Capture::default();
4987        let (endpoint, server) = spawn_test_server(
4988            Router::new()
4989                .route("/v1/messages", post(capture_ok))
4990                .with_state(capture.clone()),
4991        )
4992        .await;
4993        let state = direct_test_state(
4994            endpoint.clone(),
4995            email_pipeline(),
4996            SessionPolicy::EphemeralSingleRequest,
4997            None,
4998        );
4999
5000        let request_error = direct_proxy_inner(
5001            &state,
5002            direct_test_ingress(
5003                endpoint.clone(),
5004                DirectIngressSession::Ephemeral,
5005                WireFormat::Json,
5006                DIRECT_REQUEST_BODY,
5007            ),
5008            DIRECT_REQUEST_BODY,
5009            &LimitFailureCodec(LimitFailureDirection::Request),
5010        )
5011        .await
5012        .unwrap_err();
5013        assert_eq!(
5014            request_error.code(),
5015            ProxyErrorCode::RequestBodyLimitExceeded
5016        );
5017        assert_eq!(request_error.phase(), ProxyErrorPhase::RequestValidation);
5018        assert_eq!(request_error.http_status(), StatusCode::PAYLOAD_TOO_LARGE);
5019        assert_eq!(capture.hits.load(Ordering::SeqCst), 0);
5020        let request_response = direct_error_response(request_error);
5021        assert_eq!(request_response.status(), StatusCode::PAYLOAD_TOO_LARGE);
5022        let request_body = to_bytes(request_response.into_body(), 4 * 1024)
5023            .await
5024            .unwrap();
5025        assert_eq!(
5026            serde_json::from_slice::<serde_json::Value>(&request_body).unwrap(),
5027            serde_json::json!({
5028                "type": "error",
5029                "error": {
5030                    "type": "api_error",
5031                    "message": "proxy_validation_failed",
5032                    "code": "RequestBodyLimitExceeded",
5033                    "phase": "RequestValidation"
5034                }
5035            })
5036        );
5037
5038        let response_error = direct_proxy_inner(
5039            &state,
5040            direct_test_ingress(
5041                endpoint,
5042                DirectIngressSession::Ephemeral,
5043                WireFormat::Json,
5044                DIRECT_REQUEST_BODY,
5045            ),
5046            DIRECT_REQUEST_BODY,
5047            &LimitFailureCodec(LimitFailureDirection::Response),
5048        )
5049        .await
5050        .unwrap_err();
5051        assert_eq!(
5052            response_error.code(),
5053            ProxyErrorCode::ResponseBodyLimitExceeded
5054        );
5055        assert_eq!(response_error.phase(), ProxyErrorPhase::ResponseValidation);
5056        assert_eq!(response_error.http_status(), StatusCode::BAD_GATEWAY);
5057        assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
5058        let response = direct_error_response(response_error);
5059        assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
5060        let response_body = to_bytes(response.into_body(), 4 * 1024).await.unwrap();
5061        assert_eq!(
5062            serde_json::from_slice::<serde_json::Value>(&response_body).unwrap(),
5063            serde_json::json!({
5064                "type": "error",
5065                "error": {
5066                    "type": "api_error",
5067                    "message": "proxy_validation_failed",
5068                    "code": "ResponseBodyLimitExceeded",
5069                    "phase": "ResponseValidation"
5070                }
5071            })
5072        );
5073        server.abort();
5074    }
5075
5076    #[test]
5077    fn future_non_conflict_session_commit_failure_is_never_a_conflict() {
5078        let generation_conflict =
5079            map_session_commit_failure(SessionCommitFailureKind::GenerationConflict);
5080        assert_eq!(
5081            generation_conflict.code(),
5082            ProxyErrorCode::SessionGenerationConflict
5083        );
5084        assert_eq!(generation_conflict.phase(), ProxyErrorPhase::Session);
5085        assert_eq!(generation_conflict.http_status(), StatusCode::CONFLICT);
5086
5087        let non_conflict = map_session_commit_failure(SessionCommitFailureKind::NonConflict);
5088        assert_eq!(non_conflict.code(), ProxyErrorCode::SessionCommitFailure);
5089        assert_eq!(non_conflict.phase(), ProxyErrorPhase::Session);
5090        assert_eq!(
5091            non_conflict.http_status(),
5092            StatusCode::INTERNAL_SERVER_ERROR
5093        );
5094        assert_ne!(non_conflict.code(), generation_conflict.code());
5095    }
5096
5097    #[tokio::test]
5098    async fn continuity_lease_spans_non_stream_response_without_holding_request_mutex() {
5099        let (request, control, registry, expired_at, server) =
5100            begin_delayed_non_stream(false).await;
5101
5102        assert_eq!(control.hits.load(Ordering::SeqCst), 1);
5103        assert_eq!(registry.stats().total_leases(), 1);
5104        let concurrent_lease = registry
5105            .acquire(synthetic_principal(), NON_STREAM_SESSION_ID, None)
5106            .unwrap();
5107        let concurrent_request_guard = concurrent_lease
5108            .try_lock_request()
5109            .expect("request mutex must be free before upstream I/O");
5110        drop(concurrent_request_guard);
5111        drop(concurrent_lease);
5112        assert_eq!(registry.stats().total_leases(), 1);
5113
5114        registry.cleanup_at(expired_at);
5115        assert_eq!(registry.stats().active_entries(), 1);
5116        assert_eq!(registry.stats().tombstones(), 0);
5117
5118        control.release.notify_one();
5119        let response = request.await.unwrap().unwrap();
5120        assert_eq!(response.status(), StatusCode::OK);
5121        wait_for_lease_count(&registry, 0).await;
5122        registry.cleanup_at(expired_at);
5123        assert_eq!(registry.stats().active_entries(), 0);
5124        assert_eq!(registry.stats().tombstones(), 1);
5125        server.abort();
5126    }
5127
5128    #[tokio::test]
5129    async fn continuity_lease_releases_on_non_stream_error() {
5130        let (request, control, registry, expired_at, server) = begin_delayed_non_stream(true).await;
5131        assert_eq!(registry.stats().total_leases(), 1);
5132        registry.cleanup_at(expired_at);
5133        assert_eq!(registry.stats().active_entries(), 1);
5134
5135        control.release.notify_one();
5136        let error = request.await.unwrap().unwrap_err();
5137        assert_eq!(error.phase(), ProxyErrorPhase::ResponseValidation);
5138        assert_eq!(error.http_status(), StatusCode::BAD_GATEWAY);
5139        wait_for_lease_count(&registry, 0).await;
5140        registry.cleanup_at(expired_at);
5141        assert_eq!(registry.stats().active_entries(), 0);
5142        assert_eq!(registry.stats().tombstones(), 1);
5143        server.abort();
5144    }
5145
5146    #[tokio::test]
5147    async fn continuity_lease_releases_on_non_stream_cancellation() {
5148        let (request, _control, registry, expired_at, server) =
5149            begin_delayed_non_stream(false).await;
5150        assert_eq!(registry.stats().total_leases(), 1);
5151        registry.cleanup_at(expired_at);
5152        assert_eq!(registry.stats().active_entries(), 1);
5153
5154        request.abort();
5155        assert!(request.await.unwrap_err().is_cancelled());
5156        wait_for_lease_count(&registry, 0).await;
5157        registry.cleanup_at(expired_at);
5158        assert_eq!(registry.stats().active_entries(), 0);
5159        assert_eq!(registry.stats().tombstones(), 1);
5160        server.abort();
5161    }
5162
5163    #[tokio::test]
5164    async fn continuity_lease_spans_sse_message_stop_then_becomes_expirable() {
5165        let (response, control, registry, expired_at, server) = begin_controlled_sse(false).await;
5166        assert_eq!(control.hits.load(Ordering::SeqCst), 1);
5167        assert_eq!(registry.stats().total_leases(), 1);
5168        registry.cleanup_at(expired_at);
5169        assert_eq!(registry.stats().active_entries(), 1);
5170        assert_eq!(registry.stats().tombstones(), 0);
5171
5172        control.release.notify_one();
5173        let body = to_bytes(response.into_body(), DEFAULT_MAX_RESPONSE_BYTES)
5174            .await
5175            .unwrap();
5176        assert_eq!(body.as_ref(), DIRECT_SSE_RESPONSE);
5177        wait_for_lease_count(&registry, 0).await;
5178        registry.cleanup_at(expired_at);
5179        assert_eq!(registry.stats().active_entries(), 0);
5180        assert_eq!(registry.stats().tombstones(), 1);
5181        server.abort();
5182    }
5183
5184    #[tokio::test]
5185    async fn continuity_lease_releases_once_on_sse_error() {
5186        let (response, control, registry, expired_at, server) = begin_controlled_sse(true).await;
5187        assert_eq!(registry.stats().total_leases(), 1);
5188        registry.cleanup_at(expired_at);
5189        assert_eq!(registry.stats().active_entries(), 1);
5190
5191        control.release.notify_one();
5192        let body = to_bytes(response.into_body(), DEFAULT_MAX_RESPONSE_BYTES)
5193            .await
5194            .unwrap();
5195        assert_eq!(body.as_ref(), DIRECT_PROXY_ERROR_FRAME);
5196        wait_for_lease_count(&registry, 0).await;
5197        registry.cleanup_at(expired_at);
5198        assert_eq!(registry.stats().active_entries(), 0);
5199        assert_eq!(registry.stats().tombstones(), 1);
5200        server.abort();
5201    }
5202
5203    #[tokio::test]
5204    async fn continuity_lease_releases_once_on_sse_downstream_cancellation() {
5205        let (response, _control, registry, expired_at, server) = begin_controlled_sse(false).await;
5206        assert_eq!(registry.stats().total_leases(), 1);
5207        registry.cleanup_at(expired_at);
5208        assert_eq!(registry.stats().active_entries(), 1);
5209
5210        drop(response);
5211        wait_for_lease_count(&registry, 0).await;
5212        registry.cleanup_at(expired_at);
5213        assert_eq!(registry.stats().active_entries(), 0);
5214        assert_eq!(registry.stats().tombstones(), 1);
5215        server.abort();
5216    }
5217
5218    #[test]
5219    fn proxy_environment_is_ignored_in_isolated_helper() {
5220        let output = Command::new(std::env::current_exe().unwrap())
5221            .args([
5222                "--exact",
5223                "server::tests::proxy_environment_helper",
5224                "--nocapture",
5225            ])
5226            .env("GAZE_PROXY_ENV_HELPER", "1")
5227            .output()
5228            .unwrap();
5229        assert!(
5230            output.status.success(),
5231            "isolated proxy helper failed: {}",
5232            String::from_utf8_lossy(&output.stderr)
5233        );
5234    }
5235
5236    #[test]
5237    fn proxy_environment_helper() {
5238        if std::env::var_os("GAZE_PROXY_ENV_HELPER").is_none() {
5239            return;
5240        }
5241        tokio::runtime::Runtime::new().unwrap().block_on(async {
5242            let capture = Capture::default();
5243            let (url, server) = spawn_test_server(
5244                Router::new()
5245                    .route("/v1/messages", post(capture_ok))
5246                    .with_state(capture.clone()),
5247            )
5248            .await;
5249            let trap = TcpListener::bind("127.0.0.1:0").await.unwrap();
5250            let proxy = format!("http://{}", trap.local_addr().unwrap());
5251            for name in [
5252                "HTTP_PROXY",
5253                "HTTPS_PROXY",
5254                "ALL_PROXY",
5255                "http_proxy",
5256                "https_proxy",
5257            ] {
5258                std::env::set_var(name, &proxy);
5259            }
5260            std::env::remove_var("NO_PROXY");
5261            std::env::remove_var("no_proxy");
5262
5263            let client = DirectClient::new(DirectClientConfig::default()).unwrap();
5264            client
5265                .execute(committed_direct_request(url, WireFormat::Json, br#"{}"#))
5266                .await
5267                .unwrap();
5268            assert_eq!(capture.hits.load(Ordering::SeqCst), 1);
5269            assert!(
5270                tokio::time::timeout(Duration::from_millis(100), trap.accept())
5271                    .await
5272                    .is_err()
5273            );
5274            server.abort();
5275        });
5276    }
5277
5278    #[test]
5279    fn direct_frames_match_anthropic_codec_bytes() {
5280        assert_eq!(
5281            DIRECT_PROXY_PING_FRAME,
5282            crate::codecs::anthropic::ANTHROPIC_PROXY_PING_FRAME
5283        );
5284        assert_eq!(
5285            DIRECT_PROXY_ERROR_FRAME,
5286            crate::codecs::anthropic::ANTHROPIC_PROXY_ERROR_FRAME
5287        );
5288    }
5289
5290    #[test]
5291    fn restore_text_json_escapes_raw_values_only_into_json_syntax() {
5292        let session = Session::new(Scope::Conversation("test".to_string())).unwrap();
5293        let raw = "q\"b\\s\u{1}\u{1f}\u{7f}\u{2028}\u{e9}\r\n";
5294        let mut transaction = session.begin_transaction();
5295        let token = transaction
5296            .tokenize(&PiiClass::Custom("secret".to_string()), raw)
5297            .unwrap();
5298        transaction.commit().unwrap();
5299
5300        let none = HashSet::new();
5301        let restored = restore_text(&session, &none, &format!("{{\"v\":\"x{token}y\"}}"), true);
5302        let parsed: Value = serde_json::from_str(&restored).unwrap();
5303        assert_eq!(parsed["v"].as_str(), Some(format!("x{raw}y").as_str()));
5304        assert_eq!(
5305            restore_text(&session, &none, &format!("x{token}y"), false),
5306            format!("x{raw}y")
5307        );
5308        // A JSON-spelled manifest value that is not valid string content is escaped regardless.
5309        let json_spelled = HashSet::from([token.clone()]);
5310        let restored = restore_text(
5311            &session,
5312            &json_spelled,
5313            &format!("{{\"v\":\"{token}\"}}"),
5314            true,
5315        );
5316        assert_eq!(
5317            serde_json::from_str::<Value>(&restored).unwrap()["v"].as_str(),
5318            Some(raw)
5319        );
5320    }
5321
5322    #[test]
5323    fn json_spelled_tokens_are_the_ones_minted_inside_json_string_literals() {
5324        let mut tokens = HashSet::new();
5325        collect_json_spelled_tokens(
5326            r#"{"a":"Robert \"Bob\" Smith","b":"<Email_9>"}"#,
5327            r#"{"a":"<Custom:name_1>","<Custom:key_1>":["<Custom:name_2>"],"b":"<Email_9>"}"#,
5328            &mut tokens,
5329        );
5330        assert_eq!(
5331            tokens,
5332            HashSet::from([
5333                "<Custom:name_1>".to_string(),
5334                "<Custom:key_1>".to_string(),
5335                "<Custom:name_2>".to_string(),
5336            ])
5337        );
5338
5339        let mut none = HashSet::new();
5340        collect_json_spelled_tokens("prose", "prose <Custom:name_1>", &mut none);
5341        collect_json_spelled_tokens("{broken", "{broken <Custom:name_1>", &mut none);
5342        assert!(none.is_empty());
5343    }
5344
5345    #[test]
5346    fn restore_text_leaves_unknown_tokens() {
5347        let session = Session::new(Scope::Conversation("test".to_string())).unwrap();
5348        let none = HashSet::new();
5349        assert_eq!(
5350            restore_text(&session, &none, "hello <Email_1>", false),
5351            "hello <Email_1>"
5352        );
5353        assert_eq!(
5354            restore_text(&session, &none, "{\"to\":\"<Email_1>\"}", true),
5355            "{\"to\":\"<Email_1>\"}"
5356        );
5357    }
5358
5359    #[test]
5360    fn upstream_url_preserves_path_and_query() {
5361        let base = Url::parse("https://api.openai.com").unwrap();
5362        let uri: Uri = "/v1/chat/completions?x=1".parse().unwrap();
5363        assert_eq!(
5364            upstream_url(&base, &uri).unwrap().as_str(),
5365            "https://api.openai.com/v1/chat/completions?x=1"
5366        );
5367    }
5368
5369    #[test]
5370    fn health_lists_adapters() {
5371        let config = ProxyConfig::new(
5372            "127.0.0.1:0".parse().unwrap(),
5373            vec![Arc::new(crate::adapters::OpenAiAdapter::new(
5374                Url::parse("https://api.openai.com").unwrap(),
5375            ))],
5376        );
5377        let state = AppState {
5378            config: Arc::new(config),
5379            pipeline: Arc::new(Pipeline::builder().build().unwrap()),
5380            client: Client::new(),
5381            sessions: Arc::new(RwLock::new(HashMap::new())),
5382            direct: None,
5383            started_at: Instant::now(),
5384        };
5385        assert_eq!(health_snapshot(&state).adapters[0].name, "openai");
5386    }
5387
5388    fn span(clean: std::ops::Range<usize>, raw: std::ops::Range<usize>) -> EmittedTokenSpan {
5389        EmittedTokenSpan::new(clean, raw, gaze::PiiClass::Email)
5390    }
5391
5392    #[test]
5393    fn manifest_accounting_accepts_an_untouched_pair_and_a_described_substitution() {
5394        assert!(manifest_accounts_for_every_change(
5395            "no pii here",
5396            "no pii here",
5397            &[]
5398        ));
5399        // "lead a@b tail" -> "lead <T> tail"
5400        assert!(manifest_accounts_for_every_change(
5401            "lead a@b tail",
5402            "lead <T> tail",
5403            &[span(5..8, 5..8)],
5404        ));
5405        // Two entries, and a token longer than the bytes it replaced, so the two cursors
5406        // genuinely diverge rather than staying accidentally aligned.
5407        assert!(manifest_accounts_for_every_change(
5408            "a@b and c@d",
5409            "<TOKEN1> and <TOKEN2>",
5410            &[span(0..8, 0..3), span(13..21, 8..11)],
5411        ));
5412    }
5413
5414    #[test]
5415    fn manifest_accounting_refuses_a_deletion_no_entry_describes() {
5416        // The shape a safety-net fallback deletion produces: the bytes are gone from the clean
5417        // text and NOTHING in the manifest says where they were. Accepting this is the defect.
5418        assert!(!manifest_accounts_for_every_change(
5419            "lead deleted@example.invalid tail",
5420            "lead  tail",
5421            &[],
5422        ));
5423        // Same, alongside an unrelated entry that does reconstruct.
5424        assert!(!manifest_accounts_for_every_change(
5425            "a@b and deleted@example.invalid",
5426            "<T> and ",
5427            &[span(0..3, 0..3)],
5428        ));
5429    }
5430
5431    #[test]
5432    fn manifest_accounting_refuses_entries_that_cannot_describe_the_pair() {
5433        // Out of bounds on the raw side.
5434        assert!(!manifest_accounts_for_every_change(
5435            "short",
5436            "short",
5437            &[span(0..5, 0..99)]
5438        ));
5439        // Overlapping raw spans: the pair cannot be a sequence of disjoint substitutions.
5440        assert!(!manifest_accounts_for_every_change(
5441            "a@b c@d",
5442            "<T><T>",
5443            &[span(0..3, 0..4), span(3..6, 2..7)],
5444        ));
5445        // An entry that reconstructs, but with a tail the manifest does not account for.
5446        assert!(!manifest_accounts_for_every_change(
5447            "a@b tail",
5448            "<T> other",
5449            &[span(0..3, 0..3)]
5450        ));
5451        // Entry spans that are not char boundaries cannot index their own string.
5452        assert!(!manifest_accounts_for_every_change(
5453            "\u{00e4}x",
5454            "\u{00e4}x",
5455            &[span(0..1, 0..1)]
5456        ));
5457    }
5458
5459    #[test]
5460    fn manifest_accounting_refuses_a_deletion_that_falls_between_two_entries() {
5461        // Every other fixture reaches the fallback-deletion state with an EMPTY manifest, so the
5462        // TAIL comparison is what refuses them and the gap comparison is never the deciding
5463        // check. This is the shape where it is: two entries that both reconstruct, the last one
5464        // ending the string so both tails are empty and equal, and the deleted bytes living
5465        // entirely in the gap between them.
5466        //
5467        // Reachable, not hypothetical -- detected PII either side of a net-only marker is
5468        // ordinary provider output. Without the gap comparison this pair reads as accounted for
5469        // and `validate` admits a candidate that still holds the deleted bytes.
5470        //
5471        // raw:   "a@b GONE c@d"  entries at 0..3 and 9..12
5472        // clean: "<T1> <T2>"     entries at 0..4 and 5..9
5473        assert!(!manifest_accounts_for_every_change(
5474            "a@b GONE c@d",
5475            "<T1> <T2>",
5476            &[span(0..4, 0..3), span(5..9, 9..12)],
5477        ));
5478    }
5479}