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