agentic_core/
proxy.rs

1use std::pin::Pin;
2use std::time::Duration;
3
4use bytes::Bytes;
5use futures::{Stream, TryStreamExt};
6use http::{HeaderMap, HeaderName, HeaderValue, StatusCode};
7use reqwest::Client;
8use serde_json::Value;
9use tracing::warn;
10
11use crate::config::Config;
12use crate::error::Error;
13
14const HOP_BY_HOP: &[&str] = &[
15    "connection",
16    "keep-alive",
17    "proxy-authenticate",
18    "proxy-authorization",
19    "te",
20    "trailers",
21    "transfer-encoding",
22    "upgrade",
23];
24
25const REQUEST_DROP_EXTRA: &[&str] = &["host", "content-length"];
26
27fn is_hop_by_hop(name: &str) -> bool {
28    HOP_BY_HOP.iter().any(|h| h.eq_ignore_ascii_case(name))
29}
30
31fn is_request_drop(name: &str) -> bool {
32    is_hop_by_hop(name) || REQUEST_DROP_EXTRA.iter().any(|h| h.eq_ignore_ascii_case(name))
33}
34
35pub struct ProxyRequest {
36    pub headers: HeaderMap,
37    pub body: Bytes,
38    pub query: Option<String>,
39}
40
41pub enum ProxyBody {
42    Full(Bytes),
43    Stream(Pin<Box<dyn Stream<Item = Result<Bytes, std::io::Error>> + Send>>),
44}
45
46pub struct ProxyResponse {
47    pub status: StatusCode,
48    pub headers: HeaderMap,
49    pub body: ProxyBody,
50}
51
52#[derive(Clone)]
53pub struct ProxyState {
54    pub config: Config,
55    pub stream_client: Client,
56    pub non_stream_client: Client,
57}
58
59impl ProxyState {
60    /// # Errors
61    ///
62    /// Returns an error if the HTTP clients cannot be built.
63    pub fn new(config: Config) -> Result<Self, Error> {
64        let stream_client = Client::builder()
65            .connect_timeout(Duration::from_secs(10))
66            .timeout(Duration::from_secs(900))
67            .pool_max_idle_per_host(0)
68            .redirect(reqwest::redirect::Policy::none())
69            .build()
70            .map_err(Error::HttpClient)?;
71
72        let non_stream_client = Client::builder()
73            .connect_timeout(Duration::from_secs(10))
74            .read_timeout(Duration::from_secs(300))
75            .redirect(reqwest::redirect::Policy::none())
76            .build()
77            .map_err(Error::HttpClient)?;
78
79        Ok(Self {
80            config,
81            stream_client,
82            non_stream_client,
83        })
84    }
85}
86
87fn filter_request_headers(headers: &HeaderMap, config: &Config) -> reqwest::header::HeaderMap {
88    let mut out = reqwest::header::HeaderMap::new();
89    for (name, value) in headers {
90        if is_request_drop(name.as_str()) {
91            continue;
92        }
93        if let Ok(n) = reqwest::header::HeaderName::from_bytes(name.as_str().as_bytes()) {
94            if let Ok(v) = reqwest::header::HeaderValue::from_bytes(value.as_bytes()) {
95                out.append(n, v);
96            }
97        }
98    }
99
100    let has_auth = out.contains_key(reqwest::header::AUTHORIZATION);
101    if !has_auth {
102        if let Some(key) = config.openai_api_key.as_deref() {
103            let trimmed = key.trim();
104            if !trimmed.is_empty() {
105                if let Ok(v) = reqwest::header::HeaderValue::from_str(&format!("Bearer {trimmed}")) {
106                    out.insert(reqwest::header::AUTHORIZATION, v);
107                }
108            }
109        }
110    }
111
112    out
113}
114
115fn filter_response_headers(headers: &reqwest::header::HeaderMap) -> HeaderMap {
116    let mut out = HeaderMap::new();
117    for (name, value) in headers {
118        if is_hop_by_hop(name.as_str()) {
119            continue;
120        }
121        if let Ok(n) = HeaderName::from_bytes(name.as_str().as_bytes()) {
122            if let Ok(v) = HeaderValue::from_bytes(value.as_bytes()) {
123                out.append(n, v);
124            }
125        }
126    }
127    out
128}
129
130fn is_sse_content_type(headers: &reqwest::header::HeaderMap) -> bool {
131    headers
132        .get(reqwest::header::CONTENT_TYPE)
133        .and_then(|v| v.to_str().ok())
134        .is_some_and(|ct| ct.to_ascii_lowercase().starts_with("text/event-stream"))
135}
136
137#[must_use]
138pub fn error_response(status: StatusCode, code: &str, message: &str) -> ProxyResponse {
139    let body = serde_json::json!({
140        "error": {
141            "message": message,
142            "type": "api_error",
143            "param": null,
144            "code": code,
145        }
146    });
147    let mut headers = HeaderMap::new();
148    headers.insert("content-type", HeaderValue::from_static("application/json"));
149    ProxyResponse {
150        status,
151        headers,
152        body: ProxyBody::Full(Bytes::from(serde_json::to_vec(&body).unwrap_or_default())),
153    }
154}
155
156/// Proxy a GET request to an arbitrary upstream path.
157///
158/// Applies the same header filtering and auth injection as [`proxy_request`].
159/// Uses the non-streaming client; the response body is returned as a full
160/// [`ProxyBody::Full`] payload.
161pub async fn proxy_get(path: &str, request_headers: &HeaderMap, state: &ProxyState) -> ProxyResponse {
162    let llm_headers = filter_request_headers(request_headers, &state.config);
163    let base = state.config.llm_api_base.trim_end_matches('/');
164    let url = format!("{base}/{}", path.trim_start_matches('/'));
165
166    let llm_resp = match state.non_stream_client.get(&url).headers(llm_headers).send().await {
167        Ok(r) => r,
168        Err(e) if e.is_timeout() => {
169            warn!("upstream GET {path} timed out: {e}");
170            return error_response(StatusCode::GATEWAY_TIMEOUT, "upstream_timeout", "upstream timeout");
171        }
172        Err(e) => {
173            warn!("upstream GET {path} failed: {e}");
174            return error_response(StatusCode::BAD_GATEWAY, "upstream_unavailable", "upstream unavailable");
175        }
176    };
177
178    let status = StatusCode::from_u16(llm_resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
179    let response_headers = filter_response_headers(llm_resp.headers());
180
181    match llm_resp.bytes().await {
182        Ok(payload) => ProxyResponse {
183            status,
184            headers: response_headers,
185            body: ProxyBody::Full(payload),
186        },
187        Err(e) => {
188            warn!("failed to read upstream GET {path} body: {e}");
189            error_response(
190                StatusCode::BAD_GATEWAY,
191                "upstream_unavailable",
192                "failed to read upstream response",
193            )
194        }
195    }
196}
197
198pub async fn proxy_request(request: ProxyRequest, state: &ProxyState) -> ProxyResponse {
199    let is_streaming = serde_json::from_slice::<Value>(&request.body)
200        .ok()
201        .and_then(|v| v.get("stream")?.as_bool())
202        .unwrap_or(false);
203
204    let llm_headers = filter_request_headers(&request.headers, &state.config);
205
206    let base = state.config.llm_api_base.trim_end_matches('/');
207    let mut url = format!("{base}/v1/responses");
208    if let Some(q) = &request.query {
209        url.push('?');
210        url.push_str(q);
211    }
212
213    let client = if is_streaming {
214        &state.stream_client
215    } else {
216        &state.non_stream_client
217    };
218
219    let llm_resp = match client.post(&url).headers(llm_headers).body(request.body).send().await {
220        Ok(r) => r,
221        Err(e) if e.is_timeout() => {
222            warn!("LLM request timed out: {e}");
223            return error_response(StatusCode::GATEWAY_TIMEOUT, "llm_timeout", "LLM timeout");
224        }
225        Err(e) => {
226            warn!("LLM request failed: {e}");
227            return error_response(StatusCode::BAD_GATEWAY, "llm_unavailable", "LLM unavailable");
228        }
229    };
230
231    let status = StatusCode::from_u16(llm_resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
232    let mut response_headers = filter_response_headers(llm_resp.headers());
233
234    if is_sse_content_type(llm_resp.headers()) {
235        response_headers.insert("x-accel-buffering", HeaderValue::from_static("no"));
236
237        let byte_stream = llm_resp.bytes_stream().map_err(std::io::Error::other);
238
239        return ProxyResponse {
240            status,
241            headers: response_headers,
242            body: ProxyBody::Stream(Box::pin(byte_stream)),
243        };
244    }
245
246    let payload: Bytes = match llm_resp.bytes().await {
247        Ok(b) => b,
248        Err(e) => {
249            warn!("failed to read LLM response body: {e}");
250            return error_response(
251                StatusCode::BAD_GATEWAY,
252                "llm_unavailable",
253                "Failed to read LLM response",
254            );
255        }
256    };
257
258    ProxyResponse {
259        status,
260        headers: response_headers,
261        body: ProxyBody::Full(payload),
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use crate::config::Config;
269
270    fn test_config() -> Config {
271        Config {
272            llm_api_base: "http://localhost:8000".to_owned(),
273            openai_api_key: Some("test-key".to_owned()),
274            llm_ready_timeout_s: 5.0,
275            llm_ready_interval_s: 0.1,
276            skip_llm_ready_check: false,
277            db_url: None,
278        }
279    }
280
281    fn test_config_no_key() -> Config {
282        Config {
283            openai_api_key: None,
284            ..test_config()
285        }
286    }
287
288    #[test]
289    fn hop_by_hop_detected() {
290        assert!(is_hop_by_hop("connection"));
291        assert!(is_hop_by_hop("Connection"));
292        assert!(is_hop_by_hop("keep-alive"));
293        assert!(is_hop_by_hop("transfer-encoding"));
294        assert!(is_hop_by_hop("proxy-authorization"));
295    }
296
297    #[test]
298    fn non_hop_by_hop_passes() {
299        assert!(!is_hop_by_hop("content-type"));
300        assert!(!is_hop_by_hop("x-custom"));
301        assert!(!is_hop_by_hop("authorization"));
302    }
303
304    #[test]
305    fn request_drop_includes_host_and_content_length() {
306        assert!(is_request_drop("host"));
307        assert!(is_request_drop("content-length"));
308        assert!(is_request_drop("connection"));
309        assert!(!is_request_drop("content-type"));
310    }
311
312    #[test]
313    fn filter_request_headers_strips_hop_by_hop() {
314        let mut headers = HeaderMap::new();
315        headers.insert("content-type", "application/json".parse().unwrap());
316        headers.insert("connection", "keep-alive".parse().unwrap());
317        headers.insert("proxy-authorization", "Basic abc".parse().unwrap());
318        headers.insert("x-custom", "value".parse().unwrap());
319
320        let config = test_config_no_key();
321        let filtered = filter_request_headers(&headers, &config);
322
323        assert!(filtered.contains_key("content-type"));
324        assert!(filtered.contains_key("x-custom"));
325        assert!(!filtered.contains_key("connection"));
326        assert!(!filtered.contains_key("proxy-authorization"));
327    }
328
329    #[test]
330    fn filter_request_headers_strips_host_and_content_length() {
331        let mut headers = HeaderMap::new();
332        headers.insert("host", "example.com".parse().unwrap());
333        headers.insert("content-length", "42".parse().unwrap());
334        headers.insert("accept", "*/*".parse().unwrap());
335
336        let config = test_config_no_key();
337        let filtered = filter_request_headers(&headers, &config);
338
339        assert!(!filtered.contains_key("host"));
340        assert!(!filtered.contains_key("content-length"));
341        assert!(filtered.contains_key("accept"));
342    }
343
344    #[test]
345    fn auth_injected_when_no_client_auth() {
346        let headers = HeaderMap::new();
347        let config = test_config();
348        let filtered = filter_request_headers(&headers, &config);
349
350        assert_eq!(
351            filtered.get("authorization").unwrap().to_str().unwrap(),
352            "Bearer test-key"
353        );
354    }
355
356    #[test]
357    fn client_auth_takes_precedence() {
358        let mut headers = HeaderMap::new();
359        headers.insert("authorization", "Bearer client-token".parse().unwrap());
360
361        let config = test_config();
362        let filtered = filter_request_headers(&headers, &config);
363
364        assert_eq!(
365            filtered.get("authorization").unwrap().to_str().unwrap(),
366            "Bearer client-token"
367        );
368    }
369
370    #[test]
371    fn no_auth_injected_when_key_empty() {
372        let headers = HeaderMap::new();
373        let config = Config {
374            openai_api_key: Some("  ".to_owned()),
375            ..test_config()
376        };
377        let filtered = filter_request_headers(&headers, &config);
378
379        assert!(!filtered.contains_key("authorization"));
380    }
381
382    #[test]
383    fn no_auth_injected_when_key_none() {
384        let headers = HeaderMap::new();
385        let config = test_config_no_key();
386        let filtered = filter_request_headers(&headers, &config);
387
388        assert!(!filtered.contains_key("authorization"));
389    }
390
391    #[test]
392    fn filter_response_headers_strips_hop_by_hop() {
393        let mut headers = reqwest::header::HeaderMap::new();
394        headers.insert("content-type", "application/json".parse().unwrap());
395        headers.insert("connection", "keep-alive".parse().unwrap());
396        headers.insert("x-request-id", "abc".parse().unwrap());
397
398        let filtered = filter_response_headers(&headers);
399
400        assert!(filtered.contains_key("content-type"));
401        assert!(filtered.contains_key("x-request-id"));
402        assert!(!filtered.contains_key("connection"));
403    }
404
405    #[test]
406    fn sse_content_type_detected() {
407        let mut headers = reqwest::header::HeaderMap::new();
408        headers.insert("content-type", "text/event-stream; charset=utf-8".parse().unwrap());
409        assert!(is_sse_content_type(&headers));
410    }
411
412    #[test]
413    fn sse_content_type_case_insensitive() {
414        let mut headers = reqwest::header::HeaderMap::new();
415        headers.insert("content-type", "Text/Event-Stream".parse().unwrap());
416        assert!(is_sse_content_type(&headers));
417    }
418
419    #[test]
420    fn non_sse_content_type_rejected() {
421        let mut headers = reqwest::header::HeaderMap::new();
422        headers.insert("content-type", "application/json".parse().unwrap());
423        assert!(!is_sse_content_type(&headers));
424    }
425
426    #[test]
427    fn missing_content_type_not_sse() {
428        let headers = reqwest::header::HeaderMap::new();
429        assert!(!is_sse_content_type(&headers));
430    }
431}
agentic_core/proxy.rs

agentic_core/
proxy.rs
