lean_ctx/proxy/
forward.rs

1use axum::{
2    body::Body,
3    extract::State,
4    http::{Request, StatusCode, request::Parts},
5    response::Response,
6};
7
8use flate2::{Compression, read::GzDecoder, write::GzEncoder};
9use std::borrow::Cow;
10use std::io::{Read, Write};
11
12use super::ProxyState;
13
14/// Default request-body ceiling (MiB). A large-codebase refactor with several
15/// big files in context easily exceeds the old 10 MiB cap, which surfaced to the
16/// agent as a hard `400` mid-task. Raised and made configurable via
17/// `LEAN_CTX_PROXY_MAX_BODY_MB`.
18const DEFAULT_MAX_BODY_MB: usize = 64;
19
20pub(super) fn max_body_bytes() -> usize {
21    std::env::var("LEAN_CTX_PROXY_MAX_BODY_MB")
22        .ok()
23        .and_then(|v| v.trim().parse::<usize>().ok())
24        .filter(|mb| *mb > 0)
25        .unwrap_or(DEFAULT_MAX_BODY_MB)
26        .saturating_mul(1024 * 1024)
27}
28
29/// Transforms the already-parsed JSON request body (parsed once upstream, so the
30/// compressor never re-parses) into the serialized — possibly compressed — body,
31/// its original size, and its compressed size. A plain `fn` from the static
32/// providers or a closure that captures request-derived context (e.g. Gemini's
33/// path-encoded model) both satisfy this bound.
34pub async fn forward_request(
35    State(state): State<ProxyState>,
36    req: Request<Body>,
37    upstream_base: &str,
38    default_path: &str,
39    compress_body: impl FnOnce(serde_json::Value, usize) -> (Vec<u8>, usize, usize),
40    provider_label: &str,
41    extra_stream_types: &[&str],
42) -> Result<Response, StatusCode> {
43    let (mut parts, body) = req.into_parts();
44    let body_bytes = axum::body::to_bytes(body, max_body_bytes())
45        .await
46        .map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?;
47
48    // Org-policy gate (enterprise#25): under a signed + trusted + enforced org
49    // policy, refuse models outside the ceiling and requests over a hard
50    // budget — before any routing/compression work. No policy → no-op.
51    let gate_rules = super::policy_gate::active_rules();
52    if let Some(rules) = &gate_rules {
53        let tags = parts
54            .extensions
55            .get::<super::gateway_identity::GatewayTags>()
56            .cloned()
57            .unwrap_or_default();
58        let requested_model = requested_model_of(&parts, &body_bytes);
59        if let Err(refusal) = super::policy_gate::enforce(rules, requested_model.as_deref(), &tags)
60        {
61            tracing::warn!(
62                "lean-ctx gateway: org policy refused request ({refusal:?}) \
63                 person={:?} project={:?}",
64                tags.person,
65                tags.project
66            );
67            return Ok(super::policy_gate::refusal_response(
68                &refusal,
69                provider_label,
70            ));
71        }
72    }
73
74    // Active router (enterprise#13): may rewrite `model` in the parsed body
75    // (before compression, so exactly one serialization) and re-target the
76    // upstream within the same wire shape. Fail-open: any miss routes nothing.
77    // An org policy may exempt specific projects from downgrades (#25).
78    let routing_rules = crate::core::config::Config::load().proxy.routing.clone();
79    let downgrade_forbidden = gate_rules.as_ref().is_some_and(|rules| {
80        let project = parts
81            .extensions
82            .get::<super::gateway_identity::GatewayTags>()
83            .and_then(|t| t.project.clone());
84        super::policy_gate::downgrade_forbidden(rules, project.as_deref())
85    });
86    let route_upstreams =
87        (routing_rules.is_active() && !downgrade_forbidden).then(|| state.upstream_snapshot());
88    // Cross-shape translation (enterprise#16) only exists for the exact
89    // messages-create call — count_tokens/batches subpaths have no OpenAI
90    // equivalent and must stay within-shape.
91    let xlat_ok = cfg!(feature = "shape-xlat")
92        && provider_label == "Anthropic"
93        && parts
94            .uri
95            .path()
96            .trim_end_matches('/')
97            .ends_with("/v1/messages");
98    let route_hook = |parsed: &mut serde_json::Value| {
99        route_upstreams.as_ref().and_then(|up| {
100            super::routing::route_request(parsed, provider_label, up, &routing_rules, xlat_ok)
101        })
102    };
103
104    let prepared = prepare_request_body(
105        &parts,
106        &body_bytes,
107        compress_body,
108        route_hook,
109        upstream_base,
110        provider_label == "OpenAI",
111    )?;
112    let original_size = prepared.original_size;
113    let compressed_size = prepared.compressed_size;
114    let compression_candidate = prepared.compression_candidate;
115    let preserve_content_encoding = prepared.preserve_content_encoding;
116    let route = prepared.route;
117    let parsed = prepared.parsed;
118
119    // Apply the routing decision to the wire: re-target the upstream and — for
120    // registry providers holding their own key — swap the credential headers.
121    let upstream_base = route
122        .as_ref()
123        .and_then(|r| r.upstream_base.as_deref())
124        .unwrap_or(upstream_base);
125    if let Some(provider) = route.as_ref().and_then(|r| r.credential.as_ref()) {
126        super::providers::inject_gateway_credential(provider, &mut parts.headers)?;
127    }
128    if let Some(ref parsed) = parsed {
129        let provider = match provider_label {
130            "Anthropic" => super::introspect::Provider::Anthropic,
131            "OpenAI" | "ChatGPT" => super::introspect::Provider::OpenAi,
132            _ => super::introspect::Provider::Gemini,
133        };
134        let breakdown = super::introspect::analyze_request(parsed, provider);
135        state.introspect.record(breakdown);
136    }
137
138    // #895 Track B: assign output-savings holdout from the same pristine parsed
139    // body that each provider's compressor receives. Only when active.
140    let cohort = parsed
141        .as_ref()
142        .and_then(|p| cohort_arm(p, provider_label, default_path));
143
144    if compression_candidate {
145        state
146            .stats
147            .record_provider_request(provider_label, original_size, compressed_size);
148    }
149
150    let tokens_saved = original_size.saturating_sub(compressed_size) as u64 / 4;
151    super::metrics::record_request(tokens_saved, compressed_size as u64);
152
153    let model = parsed
154        .as_ref()
155        .and_then(|v| v.get("model"))
156        .and_then(|m| m.as_str());
157    super::cost::record(
158        model,
159        tokens_saved,
160        original_size as u64,
161        compressed_size as u64,
162    );
163
164    // Cross-shape route (enterprise#16): the body now speaks OpenAI Chat
165    // Completions — address the matching endpoint instead of the caller's
166    // `/v1/messages` path, and scan the response with the OpenAI parser.
167    let xlat = route.as_ref().is_some_and(|r| r.xlat);
168    let upstream_url = if xlat {
169        format!("{upstream_base}/v1/chat/completions")
170    } else {
171        build_upstream_url(&parts, upstream_base, default_path)
172    };
173
174    // Counterfactual probe (#701, opt-in, Anthropic native only — a
175    // cross-shape route has no Anthropic upstream to ask): fire the free
176    // count_tokens call with the ORIGINAL body, concurrent with the forward
177    // below; `usage_meter::record` reads the slot when the billed usage
178    // arrives at response end. `parsed` is the pre-compression body
179    // (compression ran on a clone) — exactly what the counterfactual must
180    // count. A detached task: it can never delay or fail the real request.
181    let counterfactual = if provider_label == "Anthropic" && !xlat {
182        super::counterfactual::maybe_spawn_probe(
183            &state.client,
184            &parts,
185            upstream_base,
186            parsed.as_ref(),
187            route.as_ref().map(|r| r.routed_from.as_str()),
188            compressed_size < original_size,
189        )
190    } else {
191        None
192    };
193
194    let response = send_upstream(
195        &state,
196        &parts,
197        &upstream_url,
198        prepared.body,
199        provider_label,
200        preserve_content_encoding,
201    )
202    .await?;
203
204    // Measured usage: read the real model + billed tokens from the response.
205    // Gemini puts the model in the URL path, not the request/response body.
206    // Translated requests get OpenAI-shape responses regardless of the label.
207    let usage_provider = if xlat {
208        super::usage::Provider::OpenAi
209    } else {
210        super::usage::Provider::from_label(provider_label)
211    };
212    let url_model = if usage_provider == super::usage::Provider::Gemini {
213        super::usage::gemini_model_from_path(parts.uri.path())
214    } else {
215        None
216    };
217
218    // Gateway context (enterprise#11/#17/#18): identity tags from the auth
219    // guard + wire savings + baseline inputs, stamped onto the usage record.
220    // A routed request is attributed to the provider actually serving it, and
221    // carries the originally requested model as routed_from (enterprise#13).
222    let mut wire = wire_context(
223        &parts,
224        provider_label,
225        upstream_base,
226        tokens_saved,
227        original_size,
228    );
229    if let Some(route) = &route {
230        wire.routed_from = Some(route.routed_from.clone());
231        if let Some(id) = &route.provider_id {
232            wire.provider = id.clone();
233        }
234        // Registry route targets carry their own local-inference flag
235        // (shadow-rate billing); built-in targets keep the URL heuristic.
236        if let Some(local) = route.local {
237            wire.is_local = local;
238        }
239    }
240    wire.counterfactual = counterfactual;
241    let wire = Some(wire);
242
243    build_response(
244        response,
245        extra_stream_types,
246        usage_provider,
247        url_model,
248        cohort,
249        wire,
250        xlat,
251    )
252    .await
253}
254
255/// Requested model for the policy gate (enterprise#25): from the JSON body
256/// (Anthropic/OpenAI dialects) or the URL path (Gemini). Encrypted-passthrough
257/// or unparseable bodies yield `None` — the ceiling governs what the gateway
258/// can see; budgets (identity-keyed) still apply to every request.
259fn requested_model_of(parts: &Parts, body_bytes: &[u8]) -> Option<String> {
260    if let Some(m) = super::usage::gemini_model_from_path(parts.uri.path()) {
261        return Some(m);
262    }
263    let decoded: Cow<'_, [u8]> = match request_body_encoding(parts) {
264        RequestBodyEncoding::Identity => Cow::Borrowed(body_bytes),
265        RequestBodyEncoding::Gzip => {
266            Cow::Owned(decode_gzip_bounded(body_bytes, max_body_bytes()).ok()?)
267        }
268        RequestBodyEncoding::Zstd => {
269            Cow::Owned(decode_zstd_bounded(body_bytes, max_body_bytes()).ok()?)
270        }
271        RequestBodyEncoding::Passthrough => return None,
272    };
273    let v: serde_json::Value = serde_json::from_slice(&decoded).ok()?;
274    v.get("model")?.as_str().map(str::to_string)
275}
276
277/// Builds the request-side [`WireContext`](super::usage::WireContext) stamped
278/// onto this turn's usage record: identity tags (inserted as a request
279/// extension by the auth guard, enterprise#11), the per-request compression
280/// saving, the pre-compression token estimate (baseline input, enterprise#18)
281/// and whether the serving upstream is local.
282fn wire_context(
283    parts: &Parts,
284    provider_label: &str,
285    upstream_base: &str,
286    tokens_saved: u64,
287    original_size: usize,
288) -> Box<super::usage::WireContext> {
289    let tags = parts
290        .extensions
291        .get::<super::gateway_identity::GatewayTags>()
292        .cloned()
293        .unwrap_or_default();
294    // Registry routes attribute usage to the provider identity ("foundry",
295    // "local"), not the wire-shape label ("OpenAI") — shape ≠ identity. The
296    // entry's resolved local flag rides along (shadow-rate billing for
297    // non-loopback local endpoints, e.g. host.docker.internal).
298    let registry = parts
299        .extensions
300        .get::<super::providers::RegistryProviderId>();
301    let provider = registry.map_or(provider_label, |r| r.id.as_str());
302    let is_local = registry.map_or_else(|| upstream_is_local(upstream_base), |r| r.local);
303    Box::new(super::usage::WireContext {
304        provider: provider.to_string(),
305        person: tags.person,
306        team: tags.team,
307        project: tags.project,
308        saved_tokens: tokens_saved,
309        // bytes/4 — the same estimation basis the proxy stats use throughout.
310        uncompressed_input_tokens: original_size as u64 / 4,
311        is_local,
312        routed_from: None,    // populated by the routing hook (wave 3)
313        counterfactual: None, // populated after the probe spawn (#701)
314    })
315}
316
317/// True when the upstream base URL points at a loopback/local endpoint (an
318/// Ollama/vLLM-style local model): billed with the transparent
319/// `local_shadow_rate` instead of provider list prices (enterprise#15/#18).
320fn upstream_is_local(upstream_base: &str) -> bool {
321    let rest = upstream_base
322        .strip_prefix("https://")
323        .or_else(|| upstream_base.strip_prefix("http://"))
324        .unwrap_or(upstream_base);
325    let host_port = rest.split(['/', '?']).next().unwrap_or(rest);
326    // Split off the port; bracketed IPv6 hosts keep their brackets.
327    let host = if let Some(b) = host_port.strip_prefix('[') {
328        b.split(']').next().unwrap_or(b)
329    } else {
330        host_port.split(':').next().unwrap_or(host_port)
331    };
332    matches!(host, "127.0.0.1" | "localhost" | "::1" | "0.0.0.0")
333}
334
335/// Output-savings arm (#895) for a request body, or `None` when no holdout is
336/// active. Keyed per provider; OpenAI's Chat vs Responses bodies are
337/// distinguished by the request path so each uses the matching cohort key.
338fn cohort_arm(
339    parsed: &serde_json::Value,
340    provider_label: &str,
341    default_path: &str,
342) -> Option<super::holdout::Arm> {
343    let holdout = crate::core::config::Config::load()
344        .proxy
345        .output_holdout_fraction();
346    if holdout <= 0.0 {
347        return None;
348    }
349    let key = match provider_label {
350        "Anthropic" => super::holdout::anthropic_key(parsed),
351        "OpenAI" | "ChatGPT" => {
352            if default_path.contains("responses") {
353                super::holdout::openai_responses_key(parsed)
354            } else {
355                super::holdout::openai_chat_key(parsed)
356            }
357        }
358        _ => super::holdout::google_key(parsed),
359    };
360    Some(super::holdout::assign(&key, holdout))
361}
362
363struct PreparedRequestBody {
364    body: Vec<u8>,
365    parsed: Option<serde_json::Value>,
366    original_size: usize,
367    compressed_size: usize,
368    compression_candidate: bool,
369    preserve_content_encoding: bool,
370    /// Routing decision applied to the body (enterprise#13); `None` = passthrough.
371    route: Option<super::routing::RouteDecision>,
372}
373
374#[derive(Clone, Copy, Debug, Eq, PartialEq)]
375enum RequestBodyEncoding {
376    Identity,
377    Gzip,
378    Zstd,
379    Passthrough,
380}
381
382fn prepare_request_body(
383    parts: &Parts,
384    body_bytes: &[u8],
385    compress_body: impl FnOnce(serde_json::Value, usize) -> (Vec<u8>, usize, usize),
386    route_hook: impl FnOnce(&mut serde_json::Value) -> Option<super::routing::RouteDecision>,
387    default_upstream_base: &str,
388    openai_shape: bool,
389) -> Result<PreparedRequestBody, StatusCode> {
390    let encoding = request_body_encoding(parts);
391    let decoded = match encoding {
392        RequestBodyEncoding::Identity => Cow::Borrowed(body_bytes),
393        RequestBodyEncoding::Gzip => Cow::Owned(decode_gzip_bounded(body_bytes, max_body_bytes())?),
394        RequestBodyEncoding::Zstd => Cow::Owned(decode_zstd_bounded(body_bytes, max_body_bytes())?),
395        RequestBodyEncoding::Passthrough => {
396            return Ok(PreparedRequestBody {
397                body: body_bytes.to_vec(),
398                parsed: None,
399                original_size: body_bytes.len(),
400                compressed_size: body_bytes.len(),
401                compression_candidate: false,
402                preserve_content_encoding: true,
403                route: None,
404            });
405        }
406    };
407
408    let Some(mut parsed) = serde_json::from_slice::<serde_json::Value>(&decoded).ok() else {
409        return Ok(PreparedRequestBody {
410            body: body_bytes.to_vec(),
411            parsed: None,
412            original_size: body_bytes.len(),
413            compressed_size: body_bytes.len(),
414            compression_candidate: false,
415            preserve_content_encoding: encoding != RequestBodyEncoding::Identity,
416            route: None,
417        });
418    };
419
420    // Router runs on the freshly parsed body, before compression: the model
421    // swap lands in the same single serialization as the compression pass.
422    let mut route = route_hook(&mut parsed);
423
424    // Measured cost opt-in (#1179): when the *effective* upstream (post-
425    // routing) is OpenRouter and the body speaks the OpenAI shape, ask for the
426    // billed charge in the final usage payload. Other upstreams never see the
427    // non-standard `usage` field (api.openai.com rejects unknown params).
428    let effective_upstream = route
429        .as_ref()
430        .and_then(|r| r.upstream_base.as_deref())
431        .unwrap_or(default_upstream_base);
432    let wants_billed_cost = super::usage_accounting::upstream_is_openrouter(effective_upstream)
433        && crate::core::config::Config::load()
434            .proxy
435            .meters_openai_usage();
436    let xlat_route = route.as_ref().is_some_and(|r| r.xlat);
437    // `usage.include` is a Chat-Completions-only parameter: gate on the shape
438    // AND the call path so a Responses-API body never carries it.
439    let chat_completions_call = openai_shape
440        && parts
441            .uri
442            .path()
443            .trim_end_matches('/')
444            .ends_with("/chat/completions");
445    if wants_billed_cost && chat_completions_call && !xlat_route {
446        super::usage_accounting::inject_usage_include(&mut parsed);
447    }
448
449    let original_size = decoded.len();
450    // Cross-shape route (enterprise#16): translate Messages→Chat-Completions
451    // and compress with the target shape's compressor. An untranslatable body
452    // fails open — the route is cancelled and the request forwards natively.
453    let (logical_body, _, compressed_size) =
454        if let Some(mut openai_body) = translated_openai_body(route.as_ref(), &parsed) {
455            if wants_billed_cost {
456                super::usage_accounting::inject_usage_include(&mut openai_body);
457            }
458            super::openai::compress_request_body(openai_body, original_size)
459        } else {
460            if route.as_ref().is_some_and(|r| r.xlat) {
461                let decision = route.take().expect("checked is_some");
462                tracing::warn!(
463                    "lean-ctx proxy: request not translatable to OpenAI shape — \
464                 cancelling route to '{}', forwarding natively",
465                    decision.provider_id.as_deref().unwrap_or("?")
466                );
467                parsed["model"] = serde_json::Value::String(decision.routed_from);
468            }
469            compress_body(parsed.clone(), original_size)
470        };
471    let body = match encoding {
472        RequestBodyEncoding::Identity => logical_body,
473        RequestBodyEncoding::Gzip => encode_gzip(&logical_body)?,
474        RequestBodyEncoding::Zstd => encode_zstd(&logical_body)?,
475        RequestBodyEncoding::Passthrough => unreachable!("passthrough returned above"),
476    };
477
478    Ok(PreparedRequestBody {
479        body,
480        parsed: Some(parsed),
481        original_size,
482        compressed_size,
483        compression_candidate: true,
484        preserve_content_encoding: encoding != RequestBodyEncoding::Identity,
485        route,
486    })
487}
488
489/// The translated OpenAI body for a cross-shape route, or `None` when the
490/// route is within-shape / absent / the body is untranslatable.
491#[cfg(feature = "shape-xlat")]
492fn translated_openai_body(
493    route: Option<&super::routing::RouteDecision>,
494    parsed: &serde_json::Value,
495) -> Option<serde_json::Value> {
496    route
497        .filter(|r| r.xlat)
498        .and_then(|_| super::shape_xlat::messages_to_chat(parsed))
499}
500
501#[cfg(not(feature = "shape-xlat"))]
502fn translated_openai_body(
503    _route: Option<&super::routing::RouteDecision>,
504    _parsed: &serde_json::Value,
505) -> Option<serde_json::Value> {
506    None
507}
508
509fn build_upstream_url(parts: &Parts, base: &str, default_path: &str) -> String {
510    format!(
511        "{base}{}",
512        parts
513            .uri
514            .path_and_query()
515            .map_or(default_path, axum::http::uri::PathAndQuery::as_str)
516    )
517}
518
519/// Request headers forwarded verbatim to the upstream provider. Anything not
520/// listed here is stripped before the request leaves the loopback proxy.
521///
522/// `openai-project` (and `openai-organization`) must be forwarded: OpenCode and
523/// the OpenAI SDK send the project scope via this header for project-scoped API
524/// keys when calling the Responses API (`/responses`). Dropping it makes OpenAI
525/// reject the request with `Missing scopes: api.responses.write` (#366).
526pub(super) const ALLOWED_REQUEST_HEADERS: &[&str] = &[
527    "authorization",
528    "x-api-key",
529    // Azure OpenAI / AI Foundry credential header (universal providers, #7).
530    "api-key",
531    "content-type",
532    "accept",
533    "user-agent",
534    "originator",
535    "anthropic-version",
536    "anthropic-beta",
537    "anthropic-dangerous-direct-browser-access",
538    "openai-organization",
539    "openai-project",
540    "openai-beta",
541    "chatgpt-account-id",
542    "x-openai-fedramp",
543    "x-openai-internal-codex-residency",
544    "x-openai-internal-codex-responses-lite",
545    "x-openai-product-sku",
546    "oai-product-sku",
547    "x-oai-attestation",
548    "x-client-request-id",
549    "x-codex-beta-features",
550    "x-codex-installation-id",
551    "x-codex-parent-thread-id",
552    "x-openai-subagent",
553    "x-codex-turn-state",
554    "x-codex-turn-metadata",
555    "x-codex-window-id",
556    "x-openai-memgen-request",
557    "x-responsesapi-include-timing-metrics",
558    "mcp-session-id",
559    "last-event-id",
560    "cache-control",
561    "x-goog-api-key",
562    "x-goog-api-client",
563];
564
565pub(super) fn is_allowed_request_header(name: &str) -> bool {
566    ALLOWED_REQUEST_HEADERS.contains(&name)
567}
568
569fn should_forward_request_header(name: &str, preserve_content_encoding: bool) -> bool {
570    is_allowed_request_header(name)
571        || (preserve_content_encoding && name.eq_ignore_ascii_case("content-encoding"))
572}
573
574fn request_body_encoding(parts: &Parts) -> RequestBodyEncoding {
575    let Some(value) = parts
576        .headers
577        .get(axum::http::header::CONTENT_ENCODING)
578        .and_then(|value| value.to_str().ok())
579    else {
580        return RequestBodyEncoding::Identity;
581    };
582
583    let encodings = value
584        .split(',')
585        .map(str::trim)
586        .filter(|part| !part.is_empty() && !part.eq_ignore_ascii_case("identity"))
587        .collect::<Vec<_>>();
588    match encodings.as_slice() {
589        [] => RequestBodyEncoding::Identity,
590        [encoding] if encoding.eq_ignore_ascii_case("gzip") => RequestBodyEncoding::Gzip,
591        [encoding] if encoding.eq_ignore_ascii_case("zstd") => RequestBodyEncoding::Zstd,
592        _ => RequestBodyEncoding::Passthrough,
593    }
594}
595
596fn decode_zstd_bounded(data: &[u8], max_bytes: usize) -> Result<Vec<u8>, StatusCode> {
597    let decoder = zstd::Decoder::new(data).map_err(|e| {
598        tracing::warn!("lean-ctx proxy: invalid zstd request body: {e}");
599        StatusCode::BAD_REQUEST
600    })?;
601    read_bounded(decoder, max_bytes).inspect_err(|e| {
602        tracing::warn!("lean-ctx proxy: zstd request decode failed: {e}");
603    })
604}
605
606fn encode_zstd(data: &[u8]) -> Result<Vec<u8>, StatusCode> {
607    zstd::encode_all(data, 3).map_err(|e| {
608        tracing::error!("lean-ctx proxy: zstd request encode failed: {e}");
609        StatusCode::INTERNAL_SERVER_ERROR
610    })
611}
612
613fn decode_gzip_bounded(data: &[u8], max_bytes: usize) -> Result<Vec<u8>, StatusCode> {
614    read_bounded(GzDecoder::new(data), max_bytes).inspect_err(|e| {
615        tracing::warn!("lean-ctx proxy: gzip request decode failed: {e}");
616    })
617}
618
619fn encode_gzip(data: &[u8]) -> Result<Vec<u8>, StatusCode> {
620    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
621    encoder.write_all(data).map_err(|e| {
622        tracing::error!("lean-ctx proxy: gzip request encode failed: {e}");
623        StatusCode::INTERNAL_SERVER_ERROR
624    })?;
625    encoder.finish().map_err(|e| {
626        tracing::error!("lean-ctx proxy: gzip request encode failed: {e}");
627        StatusCode::INTERNAL_SERVER_ERROR
628    })
629}
630
631fn read_bounded<R: Read>(reader: R, max_bytes: usize) -> Result<Vec<u8>, StatusCode> {
632    let mut limited = reader.take(max_bytes as u64 + 1);
633    let mut out = Vec::new();
634    limited
635        .read_to_end(&mut out)
636        .map_err(|_| StatusCode::BAD_REQUEST)?;
637    if out.len() > max_bytes {
638        return Err(StatusCode::PAYLOAD_TOO_LARGE);
639    }
640    Ok(out)
641}
642
643/// Statuses safe to retry once (enterprise#51): the upstream explicitly did
644/// NOT process the request (429 rejected, 502/503 gateway/unavailable). 500 and
645/// 504 are excluded — the model may have already consumed/billed the call.
646fn is_retryable_status(status: reqwest::StatusCode) -> bool {
647    matches!(status.as_u16(), 429 | 502 | 503)
648}
649
650/// Short jittered backoff before the single retry: enough for a load balancer
651/// to fail over or a rate-limit window to move, never long enough to stack up
652/// under load (fail-open rule — the client's own retry logic stays primary).
653async fn retry_backoff() {
654    let mut buf = [0u8; 2];
655    let jitter_ms =
656        getrandom::fill(&mut buf).map_or(100, |()| u64::from(u16::from_le_bytes(buf)) % 200);
657    tokio::time::sleep(std::time::Duration::from_millis(150 + jitter_ms)).await;
658}
659
660async fn send_upstream(
661    state: &ProxyState,
662    parts: &Parts,
663    url: &str,
664    body: Vec<u8>,
665    provider_label: &str,
666    preserve_content_encoding: bool,
667) -> Result<reqwest::Response, StatusCode> {
668    let send_once = |body: Vec<u8>| {
669        let mut req = state.client.request(parts.method.clone(), url);
670        for (key, value) in &parts.headers {
671            let k = key.as_str().to_lowercase();
672            if should_forward_request_header(&k, preserve_content_encoding) {
673                req = req.header(key.clone(), value.clone());
674            }
675        }
676        req.body(body).send()
677    };
678
679    // First attempt. The request body is fully buffered, and no response byte
680    // has reached the client yet — retrying here is always safe for the
681    // client connection; the status filter keeps it safe semantically.
682    let first = send_once(body.clone()).await;
683    let retry_reason = match &first {
684        Ok(resp) if is_retryable_status(resp.status()) => {
685            format!("status {}", resp.status().as_u16())
686        }
687        Err(e) if e.is_connect() || e.is_timeout() => format!("connect error: {e}"),
688        Ok(resp) => {
689            let _ = resp; // healthy (or non-retryable) response — pass through
690            return first.map_err(|_| StatusCode::BAD_GATEWAY);
691        }
692        Err(e) => {
693            tracing::error!("lean-ctx proxy: {provider_label} upstream error: {e}");
694            return Err(StatusCode::BAD_GATEWAY);
695        }
696    };
697
698    tracing::warn!("lean-ctx proxy: {provider_label} upstream {retry_reason} — retrying once");
699    retry_backoff().await;
700    match send_once(body).await {
701        Ok(resp) => Ok(resp),
702        Err(e) => {
703            // Second failure: surface the ORIGINAL outcome when it was an HTTP
704            // response (its status/headers are more honest than our 502).
705            tracing::error!("lean-ctx proxy: {provider_label} retry failed: {e}");
706            first.map_err(|_| StatusCode::BAD_GATEWAY)
707        }
708    }
709}
710
711pub(super) const FORWARDED_HEADERS: &[&str] = &[
712    "content-type",
713    "content-encoding",
714    "mcp-session-id",
715    "x-request-id",
716    "x-oai-request-id",
717    "cf-ray",
718    "x-openai-authorization-error",
719    "x-error-json",
720    "openai-organization",
721    "openai-model",
722    "openai-processing-ms",
723    "openai-version",
724    "x-models-etag",
725    "x-reasoning-included",
726    "anthropic-ratelimit-requests-limit",
727    "anthropic-ratelimit-requests-remaining",
728    "anthropic-ratelimit-tokens-limit",
729    "anthropic-ratelimit-tokens-remaining",
730    "retry-after",
731    "x-ratelimit-limit-requests",
732    "x-ratelimit-remaining-requests",
733    "x-ratelimit-limit-tokens",
734    "x-ratelimit-remaining-tokens",
735    "cache-control",
736];
737
738pub(super) fn is_forwarded_response_header(name: &str) -> bool {
739    FORWARDED_HEADERS.contains(&name)
740        || name.starts_with("x-codex-")
741        || name.starts_with("x-ratelimit-")
742}
743
744#[allow(clippy::too_many_arguments)]
745async fn build_response(
746    response: reqwest::Response,
747    extra_stream_types: &[&str],
748    usage_provider: super::usage::Provider,
749    url_model: Option<String>,
750    cohort: Option<super::holdout::Arm>,
751    wire: Option<Box<super::usage::WireContext>>,
752    xlat: bool,
753) -> Result<Response, StatusCode> {
754    let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::OK);
755    let resp_headers = response.headers().clone();
756
757    // Gateway-billed USD from response headers (#1189): LiteLLM's standard
758    // header plus the operator-configured one. Body-reported costs (OpenRouter
759    // usage.cost) beat this inside the scanner.
760    let extra_cost_header = crate::core::config::Config::load()
761        .proxy
762        .cost_response_header();
763    let header_cost =
764        super::usage_accounting::cost_from_headers(&resp_headers, extra_cost_header.as_deref());
765
766    let is_stream = resp_headers
767        .get("content-type")
768        .and_then(|v| v.to_str().ok())
769        .is_some_and(|ct| {
770            ct.contains("text/event-stream") || extra_stream_types.iter().any(|t| ct.contains(t))
771        });
772
773    if is_stream {
774        // Tee the stream through a usage Scanner: each chunk is forwarded
775        // byte-for-byte while the real model + billed tokens are extracted from
776        // the final event and recorded when the stream ends. A cross-shape
777        // route (enterprise#16) additionally translates the teed bytes back to
778        // Anthropic SSE — metering always reads the raw upstream stream.
779        let scanner = super::usage::Scanner::new(usage_provider, url_model)
780            .with_cohort(cohort)
781            .with_wire_context(wire)
782            .with_header_cost(header_cost);
783        let inner = Box::pin(response.bytes_stream());
784        let teed = Box::pin(super::usage::tee_stream(inner, scanner));
785        let body = xlat_stream_body(teed, xlat);
786        let mut resp = Response::builder().status(status);
787        for (k, v) in &resp_headers {
788            let ks = k.as_str().to_lowercase();
789            if is_forwarded_response_header(&ks) {
790                resp = resp.header(k, v);
791            }
792        }
793        return resp
794            .body(body)
795            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR);
796    }
797
798    let resp_bytes = response
799        .bytes()
800        .await
801        .map_err(|_| StatusCode::BAD_GATEWAY)?;
802
803    // Non-streaming: the whole body is one JSON object carrying `usage`.
804    let mut scanner = super::usage::Scanner::new(usage_provider, url_model)
805        .with_cohort(cohort)
806        .with_wire_context(wire)
807        .with_header_cost(header_cost);
808    scanner.feed_body(&resp_bytes);
809    if let Some(usage) = scanner.finalize() {
810        super::usage_meter::record(&usage);
811    }
812
813    let resp_bytes = if xlat {
814        xlat_response_bytes(&resp_bytes, status)
815    } else {
816        resp_bytes.to_vec()
817    };
818
819    let mut resp = Response::builder().status(status);
820    for (k, v) in &resp_headers {
821        let ks = k.as_str().to_lowercase();
822        if is_forwarded_response_header(&ks) {
823            resp = resp.header(k, v);
824        }
825    }
826    resp.body(Body::from(resp_bytes))
827        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
828}
829
830/// Streaming body, translated back to Anthropic SSE when `xlat` is set.
831#[cfg(feature = "shape-xlat")]
832fn xlat_stream_body<S>(teed: S, xlat: bool) -> Body
833where
834    S: futures::Stream<Item = Result<axum::body::Bytes, reqwest::Error>> + Send + Unpin + 'static,
835{
836    if xlat {
837        Body::from_stream(super::shape_xlat::to_anthropic_stream(teed))
838    } else {
839        Body::from_stream(teed)
840    }
841}
842
843#[cfg(not(feature = "shape-xlat"))]
844fn xlat_stream_body<S>(teed: S, _xlat: bool) -> Body
845where
846    S: futures::Stream<Item = Result<axum::body::Bytes, reqwest::Error>> + Send + Unpin + 'static,
847{
848    Body::from_stream(teed)
849}
850
851/// Non-streaming translated response: chat.completion → Anthropic message on
852/// success, error envelope on failure. Unrecognizable bodies pass unchanged
853/// (better a shape-mismatched body than a dropped one).
854#[cfg(feature = "shape-xlat")]
855fn xlat_response_bytes(resp_bytes: &[u8], status: StatusCode) -> Vec<u8> {
856    let translated = serde_json::from_slice::<serde_json::Value>(resp_bytes)
857        .ok()
858        .and_then(|v| {
859            if status.is_success() {
860                super::shape_xlat::chat_to_messages(&v)
861            } else {
862                super::shape_xlat::error_to_anthropic(&v)
863            }
864        });
865    if let Some(v) = translated {
866        serde_json::to_vec(&v).unwrap_or_else(|_| resp_bytes.to_vec())
867    } else {
868        tracing::warn!(
869            "lean-ctx proxy: cross-shape response not translatable (status {status}) — \
870             forwarding raw body"
871        );
872        resp_bytes.to_vec()
873    }
874}
875
876#[cfg(not(feature = "shape-xlat"))]
877fn xlat_response_bytes(resp_bytes: &[u8], _status: StatusCode) -> Vec<u8> {
878    resp_bytes.to_vec()
879}
880
881#[cfg(test)]
882mod tests {
883    use super::*;
884
885    fn parts_for(uri: &str) -> Parts {
886        Request::builder().uri(uri).body(()).unwrap().into_parts().0
887    }
888
889    fn add_test_marker(
890        mut value: serde_json::Value,
891        original_size: usize,
892    ) -> (Vec<u8>, usize, usize) {
893        value["lean_ctx_touched"] = serde_json::Value::Bool(true);
894        let out = serde_json::to_vec(&value).unwrap();
895        let compressed_size = out.len();
896        (out, original_size, compressed_size)
897    }
898
899    // --- enterprise#11/#18: wire context (identity + baseline inputs) ---
900
901    #[test]
902    fn upstream_is_local_detects_loopback_hosts() {
903        for local in [
904            "http://127.0.0.1:11434",
905            "http://localhost:8080/v1",
906            "http://[::1]:9999",
907            "http://0.0.0.0:4000",
908        ] {
909            assert!(upstream_is_local(local), "{local} must count as local");
910        }
911        for remote in [
912            "https://api.anthropic.com",
913            "https://acme.services.ai.azure.com/openai",
914            "https://localhost.evil.example.com", // subdomain trick ≠ local
915        ] {
916            assert!(!upstream_is_local(remote), "{remote} must not be local");
917        }
918    }
919
920    #[test]
921    fn wire_context_carries_identity_tags_and_baseline() {
922        let mut parts = parts_for("/v1/messages");
923        parts
924            .extensions
925            .insert(super::super::gateway_identity::GatewayTags {
926                person: Some("yves".into()),
927                team: Some("platform".into()),
928                project: Some("billing".into()),
929            });
930        let wire = wire_context(&parts, "Anthropic", "https://api.anthropic.com", 750, 4000);
931        assert_eq!(wire.provider, "Anthropic");
932        assert_eq!(wire.person.as_deref(), Some("yves"));
933        assert_eq!(wire.team.as_deref(), Some("platform"));
934        assert_eq!(wire.project.as_deref(), Some("billing"));
935        assert_eq!(wire.saved_tokens, 750);
936        // bytes/4 estimate, same basis as the proxy stats (enterprise#18).
937        assert_eq!(wire.uncompressed_input_tokens, 1000);
938        assert!(!wire.is_local);
939        assert_eq!(wire.routed_from, None);
940    }
941
942    #[test]
943    fn wire_context_prefers_registry_provider_id_over_shape_label() {
944        // /providers/local/... speaks the OpenAI shape but must meter as
945        // "local" — the admin breakdown groups by provider identity (#20).
946        let mut parts = parts_for("/v1/chat/completions");
947        parts
948            .extensions
949            .insert(super::super::providers::RegistryProviderId {
950                id: "local".into(),
951                local: false,
952            });
953        let wire = wire_context(&parts, "OpenAI", "http://127.0.0.1:11434", 0, 400);
954        assert_eq!(wire.provider, "local");
955    }
956
957    #[test]
958    fn wire_context_registry_local_flag_beats_url_heuristic() {
959        // The containerized gateway reaches host Ollama via
960        // host.docker.internal — not loopback, but declared local = true must
961        // book the shadow rate (enterprise#15/#18). And the inverse: a
962        // loopback-tunneled cloud endpoint declared local = false must not.
963        let mut parts = parts_for("/v1/chat/completions");
964        parts
965            .extensions
966            .insert(super::super::providers::RegistryProviderId {
967                id: "local".into(),
968                local: true,
969            });
970        let wire = wire_context(
971            &parts,
972            "OpenAI",
973            "http://host.docker.internal:11434",
974            0,
975            400,
976        );
977        assert!(wire.is_local, "declared local flag must win");
978
979        let mut parts = parts_for("/v1/chat/completions");
980        parts
981            .extensions
982            .insert(super::super::providers::RegistryProviderId {
983                id: "tunnel".into(),
984                local: false,
985            });
986        let wire = wire_context(&parts, "OpenAI", "http://127.0.0.1:9999", 0, 400);
987        assert!(!wire.is_local, "declared non-local flag must win");
988    }
989
990    // --- enterprise#51: fail-open single retry ---
991
992    #[test]
993    fn retry_covers_exactly_not_processed_statuses() {
994        // Retryable: the upstream explicitly did not process the request.
995        for code in [429_u16, 502, 503] {
996            assert!(
997                is_retryable_status(reqwest::StatusCode::from_u16(code).unwrap()),
998                "{code} must be retryable"
999            );
1000        }
1001        // Not retryable: success, client errors, and "may have processed".
1002        for code in [200_u16, 400, 401, 404, 500, 504] {
1003            assert!(
1004                !is_retryable_status(reqwest::StatusCode::from_u16(code).unwrap()),
1005                "{code} must NOT be retryable"
1006            );
1007        }
1008    }
1009
1010    #[test]
1011    fn wire_context_without_tags_still_carries_baseline() {
1012        // Local solo mode: no identity, but savings + baseline are still real.
1013        let parts = parts_for("/v1/chat/completions");
1014        let wire = wire_context(&parts, "OpenAI", "http://127.0.0.1:11434", 0, 400);
1015        assert_eq!(wire.person, None);
1016        assert_eq!(wire.project, None);
1017        assert_eq!(wire.uncompressed_input_tokens, 100);
1018        assert!(wire.is_local);
1019    }
1020
1021    #[test]
1022    fn zstd_request_bodies_are_rewritten_and_reencoded() {
1023        let body = serde_json::json!({"model": "gpt-5", "input": []});
1024        let json = serde_json::to_vec(&body).unwrap();
1025        let encoded = encode_zstd(&json).unwrap();
1026        let parts = Request::builder()
1027            .uri("/backend-api/codex/responses")
1028            .header(axum::http::header::CONTENT_ENCODING, "zstd")
1029            .body(())
1030            .unwrap()
1031            .into_parts()
1032            .0;
1033
1034        let prepared = prepare_request_body(
1035            &parts,
1036            &encoded,
1037            add_test_marker,
1038            |_| None,
1039            "https://api.openai.com",
1040            false,
1041        )
1042        .unwrap();
1043        assert_eq!(request_body_encoding(&parts), RequestBodyEncoding::Zstd);
1044        assert_eq!(prepared.original_size, json.len());
1045        assert!(prepared.compression_candidate);
1046        assert!(prepared.preserve_content_encoding);
1047        assert!(should_forward_request_header("content-encoding", true));
1048        assert!(!should_forward_request_header("content-encoding", false));
1049
1050        let decoded = zstd::decode_all(prepared.body.as_slice()).unwrap();
1051        let parsed: serde_json::Value = serde_json::from_slice(&decoded).unwrap();
1052        assert_eq!(parsed["lean_ctx_touched"], true);
1053        assert_eq!(parsed["model"], "gpt-5");
1054    }
1055
1056    #[test]
1057    fn gzip_request_bodies_are_rewritten_and_reencoded() {
1058        let body = serde_json::json!({"model": "gpt-5", "input": []});
1059        let json = serde_json::to_vec(&body).unwrap();
1060        let encoded = encode_gzip(&json).unwrap();
1061        let parts = Request::builder()
1062            .uri("/backend-api/codex/responses")
1063            .header(axum::http::header::CONTENT_ENCODING, "gzip")
1064            .body(())
1065            .unwrap()
1066            .into_parts()
1067            .0;
1068
1069        let prepared = prepare_request_body(
1070            &parts,
1071            &encoded,
1072            add_test_marker,
1073            |_| None,
1074            "https://api.openai.com",
1075            false,
1076        )
1077        .unwrap();
1078        assert_eq!(request_body_encoding(&parts), RequestBodyEncoding::Gzip);
1079        assert_eq!(prepared.original_size, json.len());
1080        assert!(prepared.compression_candidate);
1081        assert!(prepared.preserve_content_encoding);
1082
1083        let decoded = decode_gzip_bounded(&prepared.body, max_body_bytes()).unwrap();
1084        let parsed: serde_json::Value = serde_json::from_slice(&decoded).unwrap();
1085        assert_eq!(parsed["lean_ctx_touched"], true);
1086        assert_eq!(parsed["model"], "gpt-5");
1087    }
1088
1089    #[test]
1090    fn openrouter_chat_requests_opt_into_billed_cost() {
1091        let _iso = crate::core::data_dir::isolated_data_dir();
1092        let body = serde_json::json!({"model": "deepseek/deepseek-v4-flash", "messages": []});
1093        let json = serde_json::to_vec(&body).unwrap();
1094        let parts = parts_for("/v1/chat/completions");
1095
1096        let prepared = prepare_request_body(
1097            &parts,
1098            &json,
1099            add_test_marker,
1100            |_| None,
1101            "https://openrouter.ai/api",
1102            true,
1103        )
1104        .unwrap();
1105        let parsed: serde_json::Value = serde_json::from_slice(&prepared.body).unwrap();
1106        assert_eq!(
1107            parsed["usage"]["include"], true,
1108            "OpenRouter chat requests must ask for the billed cost (#1179)"
1109        );
1110    }
1111
1112    #[test]
1113    fn non_openrouter_upstreams_never_carry_the_usage_opt_in() {
1114        let _iso = crate::core::data_dir::isolated_data_dir();
1115        let body = serde_json::json!({"model": "gpt-5.5", "messages": []});
1116        let json = serde_json::to_vec(&body).unwrap();
1117        let parts = parts_for("/v1/chat/completions");
1118
1119        let prepared = prepare_request_body(
1120            &parts,
1121            &json,
1122            add_test_marker,
1123            |_| None,
1124            "https://api.openai.com",
1125            true,
1126        )
1127        .unwrap();
1128        let parsed: serde_json::Value = serde_json::from_slice(&prepared.body).unwrap();
1129        assert!(
1130            parsed.get("usage").is_none(),
1131            "api.openai.com rejects unknown top-level params — no injection"
1132        );
1133    }
1134
1135    #[test]
1136    fn responses_api_bodies_never_carry_the_usage_opt_in() {
1137        let _iso = crate::core::data_dir::isolated_data_dir();
1138        let body = serde_json::json!({"model": "gpt-5.5", "input": []});
1139        let json = serde_json::to_vec(&body).unwrap();
1140        let parts = parts_for("/v1/responses");
1141
1142        let prepared = prepare_request_body(
1143            &parts,
1144            &json,
1145            add_test_marker,
1146            |_| None,
1147            "https://openrouter.ai/api",
1148            true,
1149        )
1150        .unwrap();
1151        let parsed: serde_json::Value = serde_json::from_slice(&prepared.body).unwrap();
1152        assert!(
1153            parsed.get("usage").is_none(),
1154            "`usage.include` is Chat-Completions-only — Responses bodies stay clean"
1155        );
1156    }
1157
1158    #[test]
1159    fn identity_content_encoding_can_be_rewritten_as_json() {
1160        let parts = Request::builder()
1161            .uri("/v1/responses")
1162            .header(axum::http::header::CONTENT_ENCODING, "identity")
1163            .body(())
1164            .unwrap()
1165            .into_parts()
1166            .0;
1167
1168        assert_eq!(request_body_encoding(&parts), RequestBodyEncoding::Identity);
1169    }
1170
1171    #[test]
1172    fn unknown_encoded_request_bodies_stay_passthrough() {
1173        let parts = Request::builder()
1174            .uri("/v1/responses")
1175            .header(axum::http::header::CONTENT_ENCODING, "br")
1176            .body(())
1177            .unwrap()
1178            .into_parts()
1179            .0;
1180        let body = b"not-json";
1181
1182        let prepared = prepare_request_body(
1183            &parts,
1184            body,
1185            |_, _| panic!("unknown encodings must not be JSON-rewritten"),
1186            |_| None,
1187            "https://api.openai.com",
1188            false,
1189        )
1190        .unwrap();
1191
1192        assert_eq!(
1193            request_body_encoding(&parts),
1194            RequestBodyEncoding::Passthrough
1195        );
1196        assert_eq!(prepared.body, body);
1197        assert!(prepared.parsed.is_none());
1198        assert!(!prepared.compression_candidate);
1199        assert!(prepared.preserve_content_encoding);
1200    }
1201
1202    #[test]
1203    fn invalid_json_request_bodies_are_not_compression_candidates() {
1204        let parts = Request::builder()
1205            .uri("/v1/responses")
1206            .body(())
1207            .unwrap()
1208            .into_parts()
1209            .0;
1210        let body = b"not-json";
1211
1212        let prepared = prepare_request_body(
1213            &parts,
1214            body,
1215            |_, _| panic!("invalid JSON must not enter the compression pipeline"),
1216            |_| None,
1217            "https://api.openai.com",
1218            false,
1219        )
1220        .unwrap();
1221
1222        assert_eq!(request_body_encoding(&parts), RequestBodyEncoding::Identity);
1223        assert_eq!(prepared.body, body);
1224        assert!(prepared.parsed.is_none());
1225        assert!(!prepared.compression_candidate);
1226        assert!(!prepared.preserve_content_encoding);
1227    }
1228
1229    #[test]
1230    fn upstream_url_preserves_subpath() {
1231        let base = "https://api.anthropic.com";
1232        let parts = parts_for("/v1/messages/count_tokens");
1233        assert_eq!(
1234            build_upstream_url(&parts, base, "/v1/messages"),
1235            "https://api.anthropic.com/v1/messages/count_tokens"
1236        );
1237    }
1238
1239    #[test]
1240    fn upstream_url_preserves_batches_subpath() {
1241        let base = "https://api.anthropic.com";
1242        let parts = parts_for("/v1/messages/batches/batch_123/results");
1243        assert_eq!(
1244            build_upstream_url(&parts, base, "/v1/messages"),
1245            "https://api.anthropic.com/v1/messages/batches/batch_123/results"
1246        );
1247    }
1248
1249    #[test]
1250    fn upstream_url_exact_path() {
1251        let base = "https://api.anthropic.com";
1252        let parts = parts_for("/v1/messages");
1253        assert_eq!(
1254            build_upstream_url(&parts, base, "/v1/messages"),
1255            "https://api.anthropic.com/v1/messages"
1256        );
1257    }
1258
1259    #[test]
1260    fn upstream_url_preserves_query_params() {
1261        let base = "https://api.anthropic.com";
1262        let parts = parts_for("/v1/messages/count_tokens?model=claude-4");
1263        assert_eq!(
1264            build_upstream_url(&parts, base, "/v1/messages"),
1265            "https://api.anthropic.com/v1/messages/count_tokens?model=claude-4"
1266        );
1267    }
1268
1269    #[test]
1270    fn forwards_openai_project_and_auth_headers() {
1271        // #366: project-scoped OpenAI keys carry the scope via `OpenAI-Project`.
1272        // It must be forwarded upstream, otherwise the Responses API rejects the
1273        // call with `Missing scopes: api.responses.write`.
1274        for required in ["authorization", "openai-project", "openai-organization"] {
1275            assert!(
1276                ALLOWED_REQUEST_HEADERS.contains(&required),
1277                "request header `{required}` must be forwarded upstream"
1278            );
1279        }
1280    }
1281
1282    #[test]
1283    fn forwards_chatgpt_codex_oauth_headers() {
1284        for required in [
1285            "authorization",
1286            "chatgpt-account-id",
1287            "x-openai-fedramp",
1288            "x-openai-internal-codex-residency",
1289            "x-openai-product-sku",
1290            "oai-product-sku",
1291            "x-client-request-id",
1292            "x-codex-installation-id",
1293            "x-codex-turn-metadata",
1294            "x-openai-subagent",
1295            "x-codex-turn-state",
1296            "originator",
1297        ] {
1298            assert!(
1299                is_allowed_request_header(required),
1300                "request header `{required}` must be forwarded upstream"
1301            );
1302        }
1303    }
1304
1305    #[test]
1306    fn forwards_streamable_http_mcp_headers() {
1307        for required in ["mcp-session-id", "last-event-id"] {
1308            assert!(
1309                ALLOWED_REQUEST_HEADERS.contains(&required),
1310                "request header `{required}` must be forwarded upstream"
1311            );
1312        }
1313        assert!(
1314            is_forwarded_response_header("mcp-session-id"),
1315            "MCP session id response header must be forwarded downstream"
1316        );
1317    }
1318
1319    #[test]
1320    fn forwards_codex_state_response_headers() {
1321        for required in [
1322            "x-codex-turn-state",
1323            "x-codex-primary-used-percent",
1324            "openai-model",
1325            "x-models-etag",
1326            "x-reasoning-included",
1327            "x-oai-request-id",
1328            "cf-ray",
1329            "x-openai-authorization-error",
1330            "x-error-json",
1331        ] {
1332            assert!(
1333                is_forwarded_response_header(required),
1334                "response header `{required}` must be forwarded downstream"
1335            );
1336        }
1337    }
1338
1339    #[test]
1340    fn chatgpt_responses_use_openai_responses_holdout_key() {
1341        let _iso = crate::core::data_dir::isolated_data_dir();
1342        crate::core::config::Config::update_global(|c| {
1343            c.proxy.output_holdout = Some(1.0);
1344        })
1345        .unwrap();
1346
1347        let body = serde_json::json!({
1348            "model": "gpt-5",
1349            "input": "same conversation",
1350        });
1351
1352        assert_eq!(
1353            cohort_arm(&body, "ChatGPT", "/backend-api/codex/responses"),
1354            cohort_arm(&body, "OpenAI", "/v1/responses")
1355        );
1356    }
1357}
lean_ctx/proxy/forward.rs

lean_ctx/proxy/
forward.rs
