lean_ctx/proxy/
forward.rs

1use axum::{
2    body::Body,
3    extract::State,
4    http::{request::Parts, Request, StatusCode},
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    build_response(response, extra_stream_types).await
99}
100
101fn build_upstream_url(parts: &Parts, base: &str, default_path: &str) -> String {
102    format!(
103        "{base}{}",
104        parts
105            .uri
106            .path_and_query()
107            .map_or(default_path, axum::http::uri::PathAndQuery::as_str)
108    )
109}
110
111/// Request headers forwarded verbatim to the upstream provider. Anything not
112/// listed here is stripped before the request leaves the loopback proxy.
113///
114/// `openai-project` (and `openai-organization`) must be forwarded: OpenCode and
115/// the OpenAI SDK send the project scope via this header for project-scoped API
116/// keys when calling the Responses API (`/responses`). Dropping it makes OpenAI
117/// reject the request with `Missing scopes: api.responses.write` (#366).
118const ALLOWED_REQUEST_HEADERS: &[&str] = &[
119    "authorization",
120    "x-api-key",
121    "content-type",
122    "accept",
123    "user-agent",
124    "anthropic-version",
125    "anthropic-beta",
126    "anthropic-dangerous-direct-browser-access",
127    "openai-organization",
128    "openai-project",
129    "openai-beta",
130    "x-goog-api-key",
131    "x-goog-api-client",
132];
133
134async fn send_upstream(
135    state: &ProxyState,
136    parts: &Parts,
137    url: &str,
138    body: Vec<u8>,
139    provider_label: &str,
140) -> Result<reqwest::Response, StatusCode> {
141    let mut req = state.client.request(parts.method.clone(), url);
142
143    for (key, value) in &parts.headers {
144        let k = key.as_str().to_lowercase();
145        if ALLOWED_REQUEST_HEADERS.contains(&k.as_str()) {
146            req = req.header(key.clone(), value.clone());
147        }
148    }
149
150    req.body(body).send().await.map_err(|e| {
151        tracing::error!("lean-ctx proxy: {provider_label} upstream error: {e}");
152        StatusCode::BAD_GATEWAY
153    })
154}
155
156const FORWARDED_HEADERS: &[&str] = &[
157    "content-type",
158    "content-encoding",
159    "x-request-id",
160    "openai-organization",
161    "openai-processing-ms",
162    "openai-version",
163    "anthropic-ratelimit-requests-limit",
164    "anthropic-ratelimit-requests-remaining",
165    "anthropic-ratelimit-tokens-limit",
166    "anthropic-ratelimit-tokens-remaining",
167    "retry-after",
168    "x-ratelimit-limit-requests",
169    "x-ratelimit-remaining-requests",
170    "x-ratelimit-limit-tokens",
171    "x-ratelimit-remaining-tokens",
172    "cache-control",
173];
174
175async fn build_response(
176    response: reqwest::Response,
177    extra_stream_types: &[&str],
178) -> Result<Response, StatusCode> {
179    let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::OK);
180    let resp_headers = response.headers().clone();
181
182    let is_stream = resp_headers
183        .get("content-type")
184        .and_then(|v| v.to_str().ok())
185        .is_some_and(|ct| {
186            ct.contains("text/event-stream") || extra_stream_types.iter().any(|t| ct.contains(t))
187        });
188
189    if is_stream {
190        let stream = response.bytes_stream();
191        let body = Body::from_stream(stream);
192        let mut resp = Response::builder().status(status);
193        for (k, v) in &resp_headers {
194            let ks = k.as_str().to_lowercase();
195            if FORWARDED_HEADERS.contains(&ks.as_str()) {
196                resp = resp.header(k, v);
197            }
198        }
199        return resp
200            .body(body)
201            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR);
202    }
203
204    let resp_bytes = response
205        .bytes()
206        .await
207        .map_err(|_| StatusCode::BAD_GATEWAY)?;
208
209    let mut resp = Response::builder().status(status);
210    for (k, v) in &resp_headers {
211        let ks = k.as_str().to_lowercase();
212        if FORWARDED_HEADERS.contains(&ks.as_str()) {
213            resp = resp.header(k, v);
214        }
215    }
216    resp.body(Body::from(resp_bytes))
217        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    fn parts_for(uri: &str) -> Parts {
225        Request::builder().uri(uri).body(()).unwrap().into_parts().0
226    }
227
228    #[test]
229    fn upstream_url_preserves_subpath() {
230        let base = "https://api.anthropic.com";
231        let parts = parts_for("/v1/messages/count_tokens");
232        assert_eq!(
233            build_upstream_url(&parts, base, "/v1/messages"),
234            "https://api.anthropic.com/v1/messages/count_tokens"
235        );
236    }
237
238    #[test]
239    fn upstream_url_preserves_batches_subpath() {
240        let base = "https://api.anthropic.com";
241        let parts = parts_for("/v1/messages/batches/batch_123/results");
242        assert_eq!(
243            build_upstream_url(&parts, base, "/v1/messages"),
244            "https://api.anthropic.com/v1/messages/batches/batch_123/results"
245        );
246    }
247
248    #[test]
249    fn upstream_url_exact_path() {
250        let base = "https://api.anthropic.com";
251        let parts = parts_for("/v1/messages");
252        assert_eq!(
253            build_upstream_url(&parts, base, "/v1/messages"),
254            "https://api.anthropic.com/v1/messages"
255        );
256    }
257
258    #[test]
259    fn upstream_url_preserves_query_params() {
260        let base = "https://api.anthropic.com";
261        let parts = parts_for("/v1/messages/count_tokens?model=claude-4");
262        assert_eq!(
263            build_upstream_url(&parts, base, "/v1/messages"),
264            "https://api.anthropic.com/v1/messages/count_tokens?model=claude-4"
265        );
266    }
267
268    #[test]
269    fn forwards_openai_project_and_auth_headers() {
270        // #366: project-scoped OpenAI keys carry the scope via `OpenAI-Project`.
271        // It must be forwarded upstream, otherwise the Responses API rejects the
272        // call with `Missing scopes: api.responses.write`.
273        for required in ["authorization", "openai-project", "openai-organization"] {
274            assert!(
275                ALLOWED_REQUEST_HEADERS.contains(&required),
276                "request header `{required}` must be forwarded upstream"
277            );
278        }
279    }
280}
lean_ctx/proxy/forward.rs

lean_ctx/proxy/
forward.rs
