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