teamcraft/llm/
http.rs

1//! The one HTTP client of the LLM layer: timeouts, no redirects (a 3xx is an
2//! error, so the API key never goes to a second host), the key only in
3//! its request header.
4
5use std::time::Duration;
6
7use reqwest::header::{HeaderName, HeaderValue};
8
9use super::{AuthStyle, HttpCall, HttpResponse, LlmError, Transport};
10
11/// Time limit of one HTTP call.
12pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(180);
13
14/// HTTPS to the configured API.
15pub struct HttpTransport {
16    http: reqwest::Client,
17    base_url: String,
18    api_key: String,
19    auth: AuthStyle,
20}
21
22impl HttpTransport {
23    /// `base_url` without a trailing slash; the call's path is added.
24    pub fn new(
25        base_url: &str,
26        api_key: &str,
27        auth: AuthStyle,
28        timeout: Duration,
29    ) -> anyhow::Result<HttpTransport> {
30        let http = reqwest::Client::builder()
31            .timeout(timeout)
32            .connect_timeout(Duration::from_secs(15))
33            .redirect(reqwest::redirect::Policy::none())
34            .user_agent(concat!("TeamCraft/", env!("CARGO_PKG_VERSION")))
35            .build()?;
36        Ok(HttpTransport {
37            http,
38            base_url: base_url.trim_end_matches('/').to_string(),
39            api_key: api_key.to_string(),
40            auth,
41        })
42    }
43}
44
45impl Transport for HttpTransport {
46    async fn send(&self, call: &HttpCall) -> Result<HttpResponse, LlmError> {
47        let mut req = self
48            .http
49            .post(format!("{}{}", self.base_url, call.path))
50            .json(&call.body);
51        for (name, value) in &call.headers {
52            req = req.header(*name, value);
53        }
54        if !self.api_key.is_empty() {
55            let (name, value) = match self.auth {
56                AuthStyle::XApiKey => ("x-api-key", self.api_key.clone()),
57                AuthStyle::Bearer => ("authorization", format!("Bearer {}", self.api_key)),
58            };
59            let mut value = HeaderValue::from_str(&value)
60                .map_err(|_| LlmError::Network("the API key is not a valid header value".into()))?;
61            value.set_sensitive(true);
62            req = req.header(HeaderName::from_static(name), value);
63        }
64        let response = req.send().await.map_err(|e| {
65            if e.is_timeout() {
66                LlmError::Timeout
67            } else {
68                // without the URL: keep logs short
69                LlmError::Network(e.without_url().to_string())
70            }
71        })?;
72        let status = response.status().as_u16();
73        let headers = response.headers().clone();
74        let body = response
75            .text()
76            .await
77            .map_err(|e| LlmError::Network(e.without_url().to_string()))?;
78        Ok(HttpResponse {
79            status,
80            headers,
81            body,
82        })
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use tokio::{
89        io::{AsyncReadExt, AsyncWriteExt},
90        net::TcpListener,
91    };
92
93    use super::*;
94
95    const OK: &str = "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}";
96
97    /// Reads one request (the headers are enough) and answers `response`.
98    async fn serve_once(listener: TcpListener, response: String) -> String {
99        let (mut socket, _) = listener.accept().await.unwrap();
100        let mut buf = vec![0u8; 16 * 1024];
101        let mut seen = 0;
102        loop {
103            let n = socket.read(&mut buf[seen..]).await.unwrap();
104            seen += n;
105            if n == 0 || buf[..seen].windows(4).any(|w| w == b"\r\n\r\n") {
106                break;
107            }
108        }
109        socket.write_all(response.as_bytes()).await.unwrap();
110        socket.shutdown().await.ok();
111        String::from_utf8_lossy(&buf[..seen]).to_string()
112    }
113
114    fn call() -> HttpCall {
115        HttpCall {
116            path: "/chat/completions",
117            headers: vec![("x-extra", "1".to_string())],
118            body: serde_json::json!({ "model": "m" }),
119        }
120    }
121
122    fn transport(base: &str, key: &str, auth: AuthStyle) -> HttpTransport {
123        HttpTransport::new(base, key, auth, Duration::from_secs(5)).unwrap()
124    }
125
126    #[tokio::test]
127    async fn redirects_are_not_followed_and_the_key_stays_on_the_host() {
128        let first = TcpListener::bind("127.0.0.1:0").await.unwrap();
129        let second = TcpListener::bind("127.0.0.1:0").await.unwrap();
130        let (p1, p2) = (
131            first.local_addr().unwrap().port(),
132            second.local_addr().unwrap().port(),
133        );
134        let redirect = format!(
135            "HTTP/1.1 302 Found\r\nLocation: http://127.0.0.1:{p2}/elsewhere\r\n\
136             Content-Length: 0\r\nConnection: close\r\n\r\n"
137        );
138        let server = tokio::spawn(serve_once(first, redirect));
139        let t = transport(
140            &format!("http://127.0.0.1:{p1}/v1/"),
141            "secret-key",
142            AuthStyle::Bearer,
143        );
144        let response = t.send(&call()).await.unwrap();
145        assert_eq!(response.status, 302, "returned, not followed");
146        let request = server.await.unwrap().to_ascii_lowercase();
147        assert!(
148            request.starts_with("post /v1/chat/completions "),
149            "{request}"
150        );
151        assert!(request.contains("authorization: bearer secret-key"));
152        assert!(request.contains("x-extra: 1"));
153        // nobody connected to the second host
154        let hit = tokio::time::timeout(Duration::from_millis(300), second.accept()).await;
155        assert!(hit.is_err(), "the key must not go to another host");
156    }
157
158    #[tokio::test]
159    async fn the_key_header_follows_the_provider_and_is_optional() {
160        for (auth, expected) in [
161            (AuthStyle::XApiKey, "x-api-key: secret-key"),
162            (AuthStyle::Bearer, "authorization: bearer secret-key"),
163        ] {
164            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
165            let port = listener.local_addr().unwrap().port();
166            let server = tokio::spawn(serve_once(listener, OK.to_string()));
167            let t = transport(&format!("http://127.0.0.1:{port}"), "secret-key", auth);
168            let r = t.send(&call()).await.unwrap();
169            assert_eq!((r.status, r.body.as_str()), (200, "{}"));
170            assert!(server
171                .await
172                .unwrap()
173                .to_ascii_lowercase()
174                .contains(expected));
175        }
176        // no key (local server): no authorization header at all
177        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
178        let port = listener.local_addr().unwrap().port();
179        let server = tokio::spawn(serve_once(listener, OK.to_string()));
180        let t = transport(&format!("http://127.0.0.1:{port}"), "", AuthStyle::Bearer);
181        t.send(&call()).await.unwrap();
182        let request = server.await.unwrap().to_ascii_lowercase();
183        assert!(!request.contains("authorization") && !request.contains("x-api-key"));
184    }
185
186    #[tokio::test]
187    async fn network_errors_do_not_carry_the_url() {
188        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
189        let port = listener.local_addr().unwrap().port();
190        drop(listener);
191        let t = transport(
192            &format!("http://127.0.0.1:{port}/v1"),
193            "secret-key",
194            AuthStyle::Bearer,
195        );
196        let e = t.send(&call()).await.unwrap_err();
197        assert!(matches!(e, LlmError::Network(_)), "{e:?}");
198        let text = e.to_string();
199        assert!(
200            !text.contains("secret-key") && !text.contains("/v1"),
201            "{text}"
202        );
203    }
204}
teamcraft/llm/http.rs

teamcraft/llm/
http.rs
