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

lean_ctx/proxy/
forward.rs
