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", "accept-encoding"];
26const PROCESSED_RESPONSE_DROP_EXTRA: &[&str] = &["content-length", "content-encoding"];
27
28fn is_hop_by_hop(name: &str) -> bool {
29    HOP_BY_HOP.iter().any(|h| h.eq_ignore_ascii_case(name))
30}
31
32fn is_request_drop(name: &str) -> bool {
33    is_hop_by_hop(name) || REQUEST_DROP_EXTRA.iter().any(|h| h.eq_ignore_ascii_case(name))
34}
35
36/// Raw request data forwarded to the default Responses upstream endpoint.
37pub struct ProxyRequest {
38    pub headers: HeaderMap,
39    pub body: Bytes,
40    pub query: Option<String>,
41}
42
43/// Authentication fallback used when proxying a request to an upstream API.
44#[derive(Clone, Copy, Debug, Eq, PartialEq)]
45pub enum ProxyAuth {
46    OpenAiBearer,
47    Anthropic,
48}
49
50pub enum ProxyBody {
51    Full(Bytes),
52    Stream(Pin<Box<dyn Stream<Item = Result<Bytes, std::io::Error>> + Send>>),
53}
54
55pub struct ProxyResponse {
56    pub status: StatusCode,
57    pub headers: HeaderMap,
58    pub body: ProxyBody,
59}
60
61#[derive(Clone)]
62pub struct ProxyState {
63    pub config: Config,
64    pub stream_client: Client,
65    pub non_stream_client: Client,
66}
67
68impl ProxyState {
69    /// # Errors
70    ///
71    /// Returns an error if the HTTP clients cannot be built.
72    pub fn new(config: Config) -> Result<Self, Error> {
73        let stream_client = Client::builder()
74            .connect_timeout(Duration::from_secs(10))
75            .timeout(Duration::from_secs(900))
76            .pool_max_idle_per_host(0)
77            .redirect(reqwest::redirect::Policy::none())
78            .build()
79            .map_err(Error::HttpClient)?;
80
81        let non_stream_client = Client::builder()
82            .connect_timeout(Duration::from_secs(10))
83            .read_timeout(Duration::from_secs(300))
84            .redirect(reqwest::redirect::Policy::none())
85            .build()
86            .map_err(Error::HttpClient)?;
87
88        Ok(Self {
89            config,
90            stream_client,
91            non_stream_client,
92        })
93    }
94}
95
96/// Build the request headers forwarded to an upstream API.
97///
98/// Hop-by-hop and origin-specific headers are removed, all other headers stay
99/// open-ended, and the configured credential is injected only when the client
100/// did not supply one.
101#[must_use]
102pub fn upstream_request_headers(headers: &HeaderMap, config: &Config, auth: ProxyAuth) -> reqwest::header::HeaderMap {
103    let mut out = reqwest::header::HeaderMap::new();
104    for (name, value) in headers {
105        if is_request_drop(name.as_str()) {
106            continue;
107        }
108        if let Ok(n) = reqwest::header::HeaderName::from_bytes(name.as_str().as_bytes()) {
109            if let Ok(v) = reqwest::header::HeaderValue::from_bytes(value.as_bytes()) {
110                out.append(n, v);
111            }
112        }
113    }
114
115    let has_auth = out.contains_key(reqwest::header::AUTHORIZATION);
116    let has_api_key = out.contains_key("x-api-key");
117    if !has_auth && !has_api_key {
118        if let Some(key) = config.openai_api_key.as_deref() {
119            let trimmed = key.trim();
120            if !trimmed.is_empty() {
121                let (name, value) = match auth {
122                    ProxyAuth::OpenAiBearer => (reqwest::header::AUTHORIZATION, format!("Bearer {trimmed}")),
123                    ProxyAuth::Anthropic => (
124                        reqwest::header::HeaderName::from_static("x-api-key"),
125                        trimmed.to_owned(),
126                    ),
127                };
128                if let Ok(v) = reqwest::header::HeaderValue::from_str(&value) {
129                    out.insert(name, v);
130                }
131            }
132        }
133    }
134
135    out
136}
137
138fn filter_response_headers(headers: &reqwest::header::HeaderMap) -> HeaderMap {
139    let mut out = HeaderMap::new();
140    for (name, value) in headers {
141        if is_hop_by_hop(name.as_str()) {
142            continue;
143        }
144        if let Ok(n) = HeaderName::from_bytes(name.as_str().as_bytes()) {
145            if let Ok(v) = HeaderValue::from_bytes(value.as_bytes()) {
146                out.append(n, v);
147            }
148        }
149    }
150    out
151}
152
153/// Build response headers for an upstream body consumed or transformed in-process.
154///
155/// Hop-by-hop headers are removed along with representation metadata that may no
156/// longer describe the emitted body. Request IDs, retry guidance, and rate-limit
157/// metadata remain open-ended.
158#[must_use]
159pub fn processed_response_headers(headers: &reqwest::header::HeaderMap) -> HeaderMap {
160    let mut out = filter_response_headers(headers);
161    for name in PROCESSED_RESPONSE_DROP_EXTRA {
162        out.remove(*name);
163    }
164    out
165}
166
167fn is_sse_content_type(headers: &reqwest::header::HeaderMap) -> bool {
168    headers
169        .get(reqwest::header::CONTENT_TYPE)
170        .and_then(|v| v.to_str().ok())
171        .is_some_and(|ct| ct.to_ascii_lowercase().starts_with("text/event-stream"))
172}
173
174#[must_use]
175pub fn error_response(status: StatusCode, code: &str, message: &str) -> ProxyResponse {
176    error_response_for_auth(status, code, message, ProxyAuth::OpenAiBearer)
177}
178
179#[must_use]
180pub fn error_response_for_auth(status: StatusCode, code: &str, message: &str, auth: ProxyAuth) -> ProxyResponse {
181    let body = match auth {
182        ProxyAuth::OpenAiBearer => serde_json::json!({
183            "error": {
184                "message": message,
185                "type": "api_error",
186                "param": null,
187                "code": code,
188            }
189        }),
190        ProxyAuth::Anthropic => serde_json::json!({
191            "type": "error",
192            "error": {
193                "type": "api_error",
194                "message": message,
195            }
196        }),
197    };
198    let mut headers = HeaderMap::new();
199    headers.insert("content-type", HeaderValue::from_static("application/json"));
200    ProxyResponse {
201        status,
202        headers,
203        body: ProxyBody::Full(Bytes::from(serde_json::to_vec(&body).unwrap_or_default())),
204    }
205}
206
207/// Proxy a GET request to an arbitrary upstream path.
208///
209/// Applies the same header filtering and auth injection as [`proxy_request`].
210/// Uses the non-streaming client; the response body is returned as a full
211/// [`ProxyBody::Full`] payload.
212pub async fn proxy_get(path: &str, request_headers: &HeaderMap, state: &ProxyState) -> ProxyResponse {
213    let llm_headers = upstream_request_headers(request_headers, &state.config, ProxyAuth::OpenAiBearer);
214    let base = state.config.llm_api_base.trim_end_matches('/');
215    let url = format!("{base}/{}", path.trim_start_matches('/'));
216
217    let llm_resp = match state.non_stream_client.get(&url).headers(llm_headers).send().await {
218        Ok(r) => r,
219        Err(e) if e.is_timeout() => {
220            warn!("upstream GET {path} timed out: {e}");
221            return error_response(StatusCode::GATEWAY_TIMEOUT, "upstream_timeout", "upstream timeout");
222        }
223        Err(e) => {
224            warn!("upstream GET {path} failed: {e}");
225            return error_response(StatusCode::BAD_GATEWAY, "upstream_unavailable", "upstream unavailable");
226        }
227    };
228
229    let status = StatusCode::from_u16(llm_resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
230    let response_headers = filter_response_headers(llm_resp.headers());
231
232    match llm_resp.bytes().await {
233        Ok(payload) => ProxyResponse {
234            status,
235            headers: response_headers,
236            body: ProxyBody::Full(payload),
237        },
238        Err(e) => {
239            warn!("failed to read upstream GET {path} body: {e}");
240            error_response(
241                StatusCode::BAD_GATEWAY,
242                "upstream_unavailable",
243                "failed to read upstream response",
244            )
245        }
246    }
247}
248
249/// Proxy a request to the default `/v1/responses` upstream endpoint.
250pub async fn proxy_request(request: ProxyRequest, state: &ProxyState) -> ProxyResponse {
251    proxy_request_with_path(request, "/v1/responses", ProxyAuth::OpenAiBearer, state).await
252}
253
254/// Proxy a raw request to a selected upstream path.
255pub async fn proxy_request_with_path(
256    request: ProxyRequest,
257    path: &str,
258    auth: ProxyAuth,
259    state: &ProxyState,
260) -> ProxyResponse {
261    let is_streaming = serde_json::from_slice::<Value>(&request.body)
262        .ok()
263        .and_then(|v| v.get("stream")?.as_bool())
264        .unwrap_or(false);
265
266    let llm_headers = upstream_request_headers(&request.headers, &state.config, auth);
267
268    let base = state.config.llm_api_base.trim_end_matches('/');
269    let mut url = format!("{base}/{}", path.trim_start_matches('/'));
270    if let Some(q) = &request.query {
271        url.push('?');
272        url.push_str(q);
273    }
274
275    let client = if is_streaming {
276        &state.stream_client
277    } else {
278        &state.non_stream_client
279    };
280
281    let llm_resp = match client.post(&url).headers(llm_headers).body(request.body).send().await {
282        Ok(r) => r,
283        Err(e) if e.is_timeout() => {
284            warn!("LLM request timed out: {e}");
285            return error_response_for_auth(StatusCode::GATEWAY_TIMEOUT, "llm_timeout", "LLM timeout", auth);
286        }
287        Err(e) => {
288            warn!("LLM request failed: {e}");
289            return error_response_for_auth(StatusCode::BAD_GATEWAY, "llm_unavailable", "LLM unavailable", auth);
290        }
291    };
292
293    let status = StatusCode::from_u16(llm_resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
294    let mut response_headers = filter_response_headers(llm_resp.headers());
295
296    if is_sse_content_type(llm_resp.headers()) {
297        response_headers.insert("x-accel-buffering", HeaderValue::from_static("no"));
298
299        let byte_stream = llm_resp.bytes_stream().map_err(std::io::Error::other);
300
301        return ProxyResponse {
302            status,
303            headers: response_headers,
304            body: ProxyBody::Stream(Box::pin(byte_stream)),
305        };
306    }
307
308    let payload: Bytes = match llm_resp.bytes().await {
309        Ok(b) => b,
310        Err(e) => {
311            warn!("failed to read LLM response body: {e}");
312            return error_response_for_auth(
313                StatusCode::BAD_GATEWAY,
314                "llm_unavailable",
315                "Failed to read LLM response",
316                auth,
317            );
318        }
319    };
320
321    ProxyResponse {
322        status,
323        headers: response_headers,
324        body: ProxyBody::Full(payload),
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use crate::config::Config;
332
333    fn test_config() -> Config {
334        Config {
335            llm_api_base: "http://localhost:8000".to_owned(),
336            openai_api_key: Some("test-key".to_owned()),
337            llm_ready_timeout_s: 5.0,
338            llm_ready_interval_s: 0.1,
339            skip_llm_ready_check: false,
340            db_url: None,
341            postgres: crate::config::PostgresConfig::default(),
342            sqlite: crate::config::SqliteConfig::default(),
343            tools: crate::config::ToolRuntimeConfig::default(),
344        }
345    }
346
347    fn test_config_no_key() -> Config {
348        Config {
349            openai_api_key: None,
350            ..test_config()
351        }
352    }
353
354    #[test]
355    fn hop_by_hop_detected() {
356        assert!(is_hop_by_hop("connection"));
357        assert!(is_hop_by_hop("Connection"));
358        assert!(is_hop_by_hop("keep-alive"));
359        assert!(is_hop_by_hop("transfer-encoding"));
360        assert!(is_hop_by_hop("proxy-authorization"));
361    }
362
363    #[test]
364    fn non_hop_by_hop_passes() {
365        assert!(!is_hop_by_hop("content-type"));
366        assert!(!is_hop_by_hop("x-custom"));
367        assert!(!is_hop_by_hop("authorization"));
368    }
369
370    #[test]
371    fn request_drop_includes_host_and_content_length() {
372        assert!(is_request_drop("host"));
373        assert!(is_request_drop("content-length"));
374        assert!(is_request_drop("accept-encoding"));
375        assert!(is_request_drop("connection"));
376        assert!(!is_request_drop("content-type"));
377    }
378
379    #[test]
380    fn proxy_request_retains_legacy_construction_shape() {
381        let _request = ProxyRequest {
382            headers: HeaderMap::new(),
383            body: Bytes::new(),
384            query: None,
385        };
386    }
387
388    #[test]
389    fn filter_request_headers_strips_hop_by_hop() {
390        let mut headers = HeaderMap::new();
391        headers.insert("content-type", "application/json".parse().unwrap());
392        headers.insert("connection", "keep-alive".parse().unwrap());
393        headers.insert("proxy-authorization", "Basic abc".parse().unwrap());
394        headers.insert("x-custom", "value".parse().unwrap());
395
396        let config = test_config_no_key();
397        let filtered = upstream_request_headers(&headers, &config, ProxyAuth::OpenAiBearer);
398
399        assert!(filtered.contains_key("content-type"));
400        assert!(filtered.contains_key("x-custom"));
401        assert!(!filtered.contains_key("connection"));
402        assert!(!filtered.contains_key("proxy-authorization"));
403    }
404
405    #[test]
406    fn filter_request_headers_strips_host_and_content_length() {
407        let mut headers = HeaderMap::new();
408        headers.insert("host", "example.com".parse().unwrap());
409        headers.insert("content-length", "42".parse().unwrap());
410        headers.insert("accept", "*/*".parse().unwrap());
411
412        let config = test_config_no_key();
413        let filtered = upstream_request_headers(&headers, &config, ProxyAuth::OpenAiBearer);
414
415        assert!(!filtered.contains_key("host"));
416        assert!(!filtered.contains_key("content-length"));
417        assert!(filtered.contains_key("accept"));
418    }
419
420    #[test]
421    fn auth_injected_when_no_client_auth() {
422        let headers = HeaderMap::new();
423        let config = test_config();
424        let filtered = upstream_request_headers(&headers, &config, ProxyAuth::OpenAiBearer);
425
426        assert_eq!(
427            filtered.get("authorization").unwrap().to_str().unwrap(),
428            "Bearer test-key"
429        );
430    }
431
432    #[test]
433    fn client_auth_takes_precedence() {
434        let mut headers = HeaderMap::new();
435        headers.insert("authorization", "Bearer client-token".parse().unwrap());
436
437        let config = test_config();
438        let filtered = upstream_request_headers(&headers, &config, ProxyAuth::OpenAiBearer);
439
440        assert_eq!(
441            filtered.get("authorization").unwrap().to_str().unwrap(),
442            "Bearer client-token"
443        );
444    }
445
446    #[test]
447    fn anthropic_auth_preserves_client_api_key() {
448        let mut headers = HeaderMap::new();
449        headers.insert("x-api-key", "client-anthropic-key".parse().unwrap());
450
451        let filtered = upstream_request_headers(&headers, &test_config(), ProxyAuth::Anthropic);
452
453        assert_eq!(filtered.get("x-api-key").unwrap(), "client-anthropic-key");
454        assert!(!filtered.contains_key("authorization"));
455    }
456
457    #[test]
458    fn anthropic_auth_uses_configured_key_as_api_key_fallback() {
459        let filtered = upstream_request_headers(&HeaderMap::new(), &test_config(), ProxyAuth::Anthropic);
460
461        assert_eq!(filtered.get("x-api-key").unwrap(), "test-key");
462        assert!(!filtered.contains_key("authorization"));
463    }
464
465    #[test]
466    fn no_auth_injected_when_key_empty() {
467        let headers = HeaderMap::new();
468        let config = Config {
469            openai_api_key: Some("  ".to_owned()),
470            ..test_config()
471        };
472        let filtered = upstream_request_headers(&headers, &config, ProxyAuth::OpenAiBearer);
473
474        assert!(!filtered.contains_key("authorization"));
475    }
476
477    #[test]
478    fn no_auth_injected_when_key_none() {
479        let headers = HeaderMap::new();
480        let config = test_config_no_key();
481        let filtered = upstream_request_headers(&headers, &config, ProxyAuth::OpenAiBearer);
482
483        assert!(!filtered.contains_key("authorization"));
484    }
485
486    #[test]
487    fn filter_response_headers_strips_hop_by_hop() {
488        let mut headers = reqwest::header::HeaderMap::new();
489        headers.insert("content-type", "application/json".parse().unwrap());
490        headers.insert("connection", "keep-alive".parse().unwrap());
491        headers.insert("x-request-id", "abc".parse().unwrap());
492
493        let filtered = filter_response_headers(&headers);
494
495        assert!(filtered.contains_key("content-type"));
496        assert!(filtered.contains_key("x-request-id"));
497        assert!(!filtered.contains_key("connection"));
498    }
499
500    #[test]
501    fn processed_response_headers_preserve_metadata_and_strip_representation_headers() {
502        let mut headers = reqwest::header::HeaderMap::new();
503        headers.insert("request-id", "req_123".parse().unwrap());
504        headers.insert("retry-after", "3".parse().unwrap());
505        headers.insert("anthropic-ratelimit-requests-remaining", "7".parse().unwrap());
506        headers.insert("content-length", "99".parse().unwrap());
507        headers.insert("content-encoding", "gzip".parse().unwrap());
508
509        let filtered = processed_response_headers(&headers);
510
511        assert_eq!(filtered["request-id"], "req_123");
512        assert_eq!(filtered["retry-after"], "3");
513        assert_eq!(filtered["anthropic-ratelimit-requests-remaining"], "7");
514        assert!(!filtered.contains_key("content-length"));
515        assert!(!filtered.contains_key("content-encoding"));
516    }
517
518    #[test]
519    fn sse_content_type_detected() {
520        let mut headers = reqwest::header::HeaderMap::new();
521        headers.insert("content-type", "text/event-stream; charset=utf-8".parse().unwrap());
522        assert!(is_sse_content_type(&headers));
523    }
524
525    #[test]
526    fn sse_content_type_case_insensitive() {
527        let mut headers = reqwest::header::HeaderMap::new();
528        headers.insert("content-type", "Text/Event-Stream".parse().unwrap());
529        assert!(is_sse_content_type(&headers));
530    }
531
532    #[test]
533    fn non_sse_content_type_rejected() {
534        let mut headers = reqwest::header::HeaderMap::new();
535        headers.insert("content-type", "application/json".parse().unwrap());
536        assert!(!is_sse_content_type(&headers));
537    }
538
539    #[test]
540    fn missing_content_type_not_sse() {
541        let headers = reqwest::header::HeaderMap::new();
542        assert!(!is_sse_content_type(&headers));
543    }
544}
agentic_core/proxy.rs

agentic_core/
proxy.rs
