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

lean_ctx/proxy/
forward.rs
