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 super::ProxyState;
9
10/// Default request-body ceiling (MiB). A large-codebase refactor with several
11/// big files in context easily exceeds the old 10 MiB cap, which surfaced to the
12/// agent as a hard `400` mid-task. Raised and made configurable via
13/// `LEAN_CTX_PROXY_MAX_BODY_MB`.
14const DEFAULT_MAX_BODY_MB: usize = 64;
15
16fn max_body_bytes() -> usize {
17    std::env::var("LEAN_CTX_PROXY_MAX_BODY_MB")
18        .ok()
19        .and_then(|v| v.trim().parse::<usize>().ok())
20        .filter(|mb| *mb > 0)
21        .unwrap_or(DEFAULT_MAX_BODY_MB)
22        .saturating_mul(1024 * 1024)
23}
24
25/// Receives the already-parsed JSON value, avoiding a redundant
26/// `serde_json::from_slice` on every request. Returns the serialized (possibly
27/// compressed) body, original size, and compressed size.
28pub type CompressFn = fn(serde_json::Value, usize) -> (Vec<u8>, usize, usize);
29
30pub async fn forward_request(
31    State(state): State<ProxyState>,
32    req: Request<Body>,
33    upstream_base: &str,
34    default_path: &str,
35    compress_body: CompressFn,
36    provider_label: &str,
37    extra_stream_types: &[&str],
38) -> Result<Response, StatusCode> {
39    let (parts, body) = req.into_parts();
40    let body_bytes = axum::body::to_bytes(body, max_body_bytes())
41        .await
42        .map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?;
43
44    state.stats.record_request();
45
46    let original_size = body_bytes.len();
47
48    // Parse once; the parsed value is shared between introspection, cost
49    // attribution, and compression — eliminating the redundant re-parse that
50    // each compress_body function previously performed internally.
51    let parsed = serde_json::from_slice::<serde_json::Value>(&body_bytes).ok();
52    if let Some(ref parsed) = parsed {
53        let provider = match provider_label {
54            "Anthropic" => super::introspect::Provider::Anthropic,
55            "OpenAI" => super::introspect::Provider::OpenAi,
56            _ => super::introspect::Provider::Gemini,
57        };
58        let breakdown = super::introspect::analyze_request(parsed, provider);
59        state.introspect.record(breakdown);
60    }
61
62    let (compressed_body, _, compressed_size) = if let Some(value) = parsed.clone() {
63        compress_body(value, original_size)
64    } else {
65        (body_bytes.to_vec(), original_size, original_size)
66    };
67
68    if compressed_size < original_size {
69        state
70            .stats
71            .record_compression(original_size, compressed_size);
72    }
73
74    let tokens_saved = original_size.saturating_sub(compressed_size) as u64 / 4;
75    super::metrics::record_request(tokens_saved, compressed_size as u64);
76
77    let model = parsed
78        .as_ref()
79        .and_then(|v| v.get("model"))
80        .and_then(|m| m.as_str());
81    super::cost::record(
82        model,
83        tokens_saved,
84        original_size as u64,
85        compressed_size as u64,
86    );
87
88    let upstream_url = build_upstream_url(&parts, upstream_base, default_path);
89    let response = send_upstream(
90        &state,
91        &parts,
92        &upstream_url,
93        compressed_body,
94        provider_label,
95    )
96    .await?;
97
98    // Measured usage: read the real model + billed tokens from the response.
99    // Gemini puts the model in the URL path, not the request/response body.
100    let usage_provider = super::usage::Provider::from_label(provider_label);
101    let url_model = if usage_provider == super::usage::Provider::Gemini {
102        super::usage::gemini_model_from_path(parts.uri.path())
103    } else {
104        None
105    };
106
107    build_response(response, extra_stream_types, usage_provider, url_model).await
108}
109
110fn build_upstream_url(parts: &Parts, base: &str, default_path: &str) -> String {
111    format!(
112        "{base}{}",
113        parts
114            .uri
115            .path_and_query()
116            .map_or(default_path, axum::http::uri::PathAndQuery::as_str)
117    )
118}
119
120/// Request headers forwarded verbatim to the upstream provider. Anything not
121/// listed here is stripped before the request leaves the loopback proxy.
122///
123/// `openai-project` (and `openai-organization`) must be forwarded: OpenCode and
124/// the OpenAI SDK send the project scope via this header for project-scoped API
125/// keys when calling the Responses API (`/responses`). Dropping it makes OpenAI
126/// reject the request with `Missing scopes: api.responses.write` (#366).
127const ALLOWED_REQUEST_HEADERS: &[&str] = &[
128    "authorization",
129    "x-api-key",
130    "content-type",
131    "accept",
132    "user-agent",
133    "anthropic-version",
134    "anthropic-beta",
135    "anthropic-dangerous-direct-browser-access",
136    "openai-organization",
137    "openai-project",
138    "openai-beta",
139    "x-goog-api-key",
140    "x-goog-api-client",
141];
142
143async fn send_upstream(
144    state: &ProxyState,
145    parts: &Parts,
146    url: &str,
147    body: Vec<u8>,
148    provider_label: &str,
149) -> Result<reqwest::Response, StatusCode> {
150    let mut req = state.client.request(parts.method.clone(), url);
151
152    for (key, value) in &parts.headers {
153        let k = key.as_str().to_lowercase();
154        if ALLOWED_REQUEST_HEADERS.contains(&k.as_str()) {
155            req = req.header(key.clone(), value.clone());
156        }
157    }
158
159    req.body(body).send().await.map_err(|e| {
160        tracing::error!("lean-ctx proxy: {provider_label} upstream error: {e}");
161        StatusCode::BAD_GATEWAY
162    })
163}
164
165const FORWARDED_HEADERS: &[&str] = &[
166    "content-type",
167    "content-encoding",
168    "x-request-id",
169    "openai-organization",
170    "openai-processing-ms",
171    "openai-version",
172    "anthropic-ratelimit-requests-limit",
173    "anthropic-ratelimit-requests-remaining",
174    "anthropic-ratelimit-tokens-limit",
175    "anthropic-ratelimit-tokens-remaining",
176    "retry-after",
177    "x-ratelimit-limit-requests",
178    "x-ratelimit-remaining-requests",
179    "x-ratelimit-limit-tokens",
180    "x-ratelimit-remaining-tokens",
181    "cache-control",
182];
183
184async fn build_response(
185    response: reqwest::Response,
186    extra_stream_types: &[&str],
187    usage_provider: super::usage::Provider,
188    url_model: Option<String>,
189) -> Result<Response, StatusCode> {
190    let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::OK);
191    let resp_headers = response.headers().clone();
192
193    let is_stream = resp_headers
194        .get("content-type")
195        .and_then(|v| v.to_str().ok())
196        .is_some_and(|ct| {
197            ct.contains("text/event-stream") || extra_stream_types.iter().any(|t| ct.contains(t))
198        });
199
200    if is_stream {
201        // Tee the stream through a usage Scanner: each chunk is forwarded
202        // byte-for-byte while the real model + billed tokens are extracted from
203        // the final event and recorded when the stream ends.
204        let scanner = super::usage::Scanner::new(usage_provider, url_model);
205        let inner = Box::pin(response.bytes_stream());
206        let body = Body::from_stream(super::usage::tee_stream(inner, scanner));
207        let mut resp = Response::builder().status(status);
208        for (k, v) in &resp_headers {
209            let ks = k.as_str().to_lowercase();
210            if FORWARDED_HEADERS.contains(&ks.as_str()) {
211                resp = resp.header(k, v);
212            }
213        }
214        return resp
215            .body(body)
216            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR);
217    }
218
219    let resp_bytes = response
220        .bytes()
221        .await
222        .map_err(|_| StatusCode::BAD_GATEWAY)?;
223
224    // Non-streaming: the whole body is one JSON object carrying `usage`.
225    let mut scanner = super::usage::Scanner::new(usage_provider, url_model);
226    scanner.feed_body(&resp_bytes);
227    if let Some(usage) = scanner.finalize() {
228        super::usage_meter::record(&usage);
229    }
230
231    let mut resp = Response::builder().status(status);
232    for (k, v) in &resp_headers {
233        let ks = k.as_str().to_lowercase();
234        if FORWARDED_HEADERS.contains(&ks.as_str()) {
235            resp = resp.header(k, v);
236        }
237    }
238    resp.body(Body::from(resp_bytes))
239        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    fn parts_for(uri: &str) -> Parts {
247        Request::builder().uri(uri).body(()).unwrap().into_parts().0
248    }
249
250    #[test]
251    fn upstream_url_preserves_subpath() {
252        let base = "https://api.anthropic.com";
253        let parts = parts_for("/v1/messages/count_tokens");
254        assert_eq!(
255            build_upstream_url(&parts, base, "/v1/messages"),
256            "https://api.anthropic.com/v1/messages/count_tokens"
257        );
258    }
259
260    #[test]
261    fn upstream_url_preserves_batches_subpath() {
262        let base = "https://api.anthropic.com";
263        let parts = parts_for("/v1/messages/batches/batch_123/results");
264        assert_eq!(
265            build_upstream_url(&parts, base, "/v1/messages"),
266            "https://api.anthropic.com/v1/messages/batches/batch_123/results"
267        );
268    }
269
270    #[test]
271    fn upstream_url_exact_path() {
272        let base = "https://api.anthropic.com";
273        let parts = parts_for("/v1/messages");
274        assert_eq!(
275            build_upstream_url(&parts, base, "/v1/messages"),
276            "https://api.anthropic.com/v1/messages"
277        );
278    }
279
280    #[test]
281    fn upstream_url_preserves_query_params() {
282        let base = "https://api.anthropic.com";
283        let parts = parts_for("/v1/messages/count_tokens?model=claude-4");
284        assert_eq!(
285            build_upstream_url(&parts, base, "/v1/messages"),
286            "https://api.anthropic.com/v1/messages/count_tokens?model=claude-4"
287        );
288    }
289
290    #[test]
291    fn forwards_openai_project_and_auth_headers() {
292        // #366: project-scoped OpenAI keys carry the scope via `OpenAI-Project`.
293        // It must be forwarded upstream, otherwise the Responses API rejects the
294        // call with `Missing scopes: api.responses.write`.
295        for required in ["authorization", "openai-project", "openai-organization"] {
296            assert!(
297                ALLOWED_REQUEST_HEADERS.contains(&required),
298                "request header `{required}` must be forwarded upstream"
299            );
300        }
301    }
302}
lean_ctx/proxy/forward.rs

lean_ctx/proxy/
forward.rs
