tenant.openagents/omega

1use std::error::Error;
2use std::sync::{LazyLock, OnceLock};
3use std::{borrow::Cow, mem, pin::Pin, task::Poll, time::Duration};
4
5use gpui_util::defer;
6
7use anyhow::anyhow;
8use bytes::{BufMut, Bytes, BytesMut};
9use futures::{AsyncRead, FutureExt as _, TryStreamExt as _};
10use http_client::{RedirectPolicy, Url, http};
11use regex::Regex;
12use reqwest::{
13    header::{HeaderMap, HeaderValue},
14    redirect,
15};
16
17const DEFAULT_CAPACITY: usize = 4096;
18static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
19static REDACT_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"key=[^&]+").unwrap());
20
21pub struct ReqwestClient {
22    client: reqwest::Client,
23    proxy: Option<Url>,
24    user_agent: Option<HeaderValue>,
25    handle: tokio::runtime::Handle,
26}
27
28impl ReqwestClient {
29    fn builder() -> reqwest::ClientBuilder {
30        reqwest::Client::builder()
31            .use_rustls_tls()
32            .connect_timeout(Duration::from_secs(10))
33            // Detect and drop connections that have silently gone bad on a
34            // flaky path (NAT timeouts, resets) instead of reusing them. A
35            // stale reused HTTP/2 connection is a common source of
36            // `BadRecordMac` TLS errors against long-lived endpoints.
37            .tcp_keepalive(Duration::from_secs(30))
38            .pool_idle_timeout(Duration::from_secs(30))
39            .http2_keep_alive_interval(Duration::from_secs(15))
40            .http2_keep_alive_timeout(Duration::from_secs(10))
41            .http2_keep_alive_while_idle(true)
42    }
43
44    pub fn new() -> Self {
45        Self::builder()
46            .build()
47            .expect("Failed to initialize HTTP client")
48            .into()
49    }
50
51    pub fn user_agent(agent: &str) -> anyhow::Result<Self> {
52        let mut map = HeaderMap::new();
53        map.insert(http::header::USER_AGENT, HeaderValue::from_str(agent)?);
54        let client = Self::builder().default_headers(map).build()?;
55        Ok(client.into())
56    }
57
58    pub fn proxy_and_user_agent(proxy: Option<Url>, user_agent: &str) -> anyhow::Result<Self> {
59        let user_agent = HeaderValue::from_str(user_agent)?;
60
61        let mut map = HeaderMap::new();
62        map.insert(http::header::USER_AGENT, user_agent.clone());
63        let mut client = Self::builder().default_headers(map);
64        let client_has_proxy;
65
66        if let Some(proxy) = proxy.as_ref().and_then(|proxy_url| {
67            reqwest::Proxy::all(proxy_url.clone())
68                .inspect_err(|e| {
69                    log::error!(
70                        "Failed to parse proxy URL '{}': {}",
71                        proxy_url,
72                        e.source().unwrap_or(&e as &_)
73                    )
74                })
75                .ok()
76        }) {
77            // Respect NO_PROXY env var
78            client = client.proxy(proxy.no_proxy(reqwest::NoProxy::from_env()));
79            client_has_proxy = true;
80        } else {
81            client_has_proxy = false;
82        };
83
84        let client = client
85            .use_preconfigured_tls(http_client_tls::tls_config())
86            .build()?;
87        let mut client: ReqwestClient = client.into();
88        client.proxy = client_has_proxy.then_some(proxy).flatten();
89        client.user_agent = Some(user_agent);
90        Ok(client)
91    }
92}
93
94pub fn runtime() -> &'static tokio::runtime::Runtime {
95    RUNTIME.get_or_init(|| {
96        tokio::runtime::Builder::new_multi_thread()
97            // Since we now have two executors, let's try to keep our footprint small
98            .worker_threads(1)
99            .enable_all()
100            .build()
101            .expect("Failed to initialize HTTP client")
102    })
103}
104
105impl From<reqwest::Client> for ReqwestClient {
106    fn from(client: reqwest::Client) -> Self {
107        let handle = tokio::runtime::Handle::try_current().unwrap_or_else(|_| {
108            log::debug!("no tokio runtime found, creating one for Reqwest...");
109            runtime().handle().clone()
110        });
111        Self {
112            client,
113            handle,
114            proxy: None,
115            user_agent: None,
116        }
117    }
118}
119
120// This struct is essentially a re-implementation of
121// https://docs.rs/tokio-util/0.7.12/tokio_util/io/struct.ReaderStream.html
122// except outside of Tokio's aegis
123struct StreamReader {
124    reader: Option<Pin<Box<dyn futures::AsyncRead + Send + Sync>>>,
125    buf: BytesMut,
126    capacity: usize,
127}
128
129impl StreamReader {
130    fn new(reader: Pin<Box<dyn futures::AsyncRead + Send + Sync>>) -> Self {
131        Self {
132            reader: Some(reader),
133            buf: BytesMut::new(),
134            capacity: DEFAULT_CAPACITY,
135        }
136    }
137}
138
139impl futures::Stream for StreamReader {
140    type Item = std::io::Result<Bytes>;
141
142    fn poll_next(
143        mut self: Pin<&mut Self>,
144        cx: &mut std::task::Context<'_>,
145    ) -> Poll<Option<Self::Item>> {
146        let mut this = self.as_mut();
147
148        let mut reader = match this.reader.take() {
149            Some(r) => r,
150            None => return Poll::Ready(None),
151        };
152
153        if this.buf.capacity() == 0 {
154            let capacity = this.capacity;
155            this.buf.reserve(capacity);
156        }
157
158        match poll_read_buf(&mut reader, cx, &mut this.buf) {
159            Poll::Pending => {
160                self.reader = Some(reader);
161
162                Poll::Pending
163            }
164            Poll::Ready(Err(err)) => {
165                self.reader = None;
166
167                Poll::Ready(Some(Err(err)))
168            }
169            Poll::Ready(Ok(0)) => {
170                self.reader = None;
171                Poll::Ready(None)
172            }
173            Poll::Ready(Ok(_)) => {
174                let chunk = this.buf.split();
175                self.reader = Some(reader);
176                Poll::Ready(Some(Ok(chunk.freeze())))
177            }
178        }
179    }
180}
181
182/// Implementation from <https://docs.rs/tokio-util/0.7.12/src/tokio_util/util/poll_buf.rs.html>
183/// Specialized for this use case
184fn poll_read_buf(
185    io: &mut Pin<Box<dyn futures::AsyncRead + Send + Sync>>,
186    cx: &mut std::task::Context<'_>,
187    buf: &mut BytesMut,
188) -> Poll<std::io::Result<usize>> {
189    if !buf.has_remaining_mut() {
190        return Poll::Ready(Ok(0));
191    }
192
193    let n = {
194        let dst = buf.chunk_mut();
195
196        // Safety: `chunk_mut()` returns a `&mut UninitSlice`, and `UninitSlice` is a
197        // transparent wrapper around `[std::mem::MaybeUninit<u8>]`.
198        let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>]) };
199        let mut read_buf = tokio::io::ReadBuf::uninit(dst);
200        let unfilled_portion = read_buf.initialize_unfilled();
201        // SAFETY: Pin projection
202        let io_pin = unsafe { Pin::new_unchecked(io) };
203        // `futures::AsyncRead` reports the byte count as the poll's return
204        // value; `read_buf.filled()` stays empty because the reader writes
205        // through the initialized slice without advancing the `ReadBuf`.
206        std::task::ready!(io_pin.poll_read(cx, unfilled_portion)?)
207    };
208
209    // Safety: `initialize_unfilled()` zero-initialized the entire spare
210    // capacity, so the first `n` bytes are initialized no matter how many the
211    // reader actually wrote, and `advance_mut` panics rather than exceeding
212    // the capacity if `n` overstates the slice length.
213    unsafe {
214        buf.advance_mut(n);
215    }
216
217    Poll::Ready(Ok(n))
218}
219
220fn redact_error(mut error: reqwest::Error) -> reqwest::Error {
221    if let Some(url) = error.url_mut()
222        && let Some(query) = url.query()
223        && let Cow::Owned(redacted) = REDACT_REGEX.replace_all(query, "key=REDACTED")
224    {
225        url.set_query(Some(redacted.as_str()));
226    }
227    error
228}
229
230impl http_client::HttpClient for ReqwestClient {
231    fn proxy(&self) -> Option<&Url> {
232        self.proxy.as_ref()
233    }
234
235    fn user_agent(&self) -> Option<&HeaderValue> {
236        self.user_agent.as_ref()
237    }
238
239    fn send(
240        &self,
241        req: http::Request<http_client::AsyncBody>,
242    ) -> futures::future::BoxFuture<
243        'static,
244        anyhow::Result<http_client::Response<http_client::AsyncBody>>,
245    > {
246        let (parts, body) = req.into_parts();
247
248        let mut request = self.client.request(parts.method, parts.uri.to_string());
249        request = request.headers(parts.headers);
250        if let Some(redirect_policy) = parts.extensions.get::<RedirectPolicy>() {
251            request = request.redirect_policy(match redirect_policy {
252                RedirectPolicy::NoFollow => redirect::Policy::none(),
253                RedirectPolicy::FollowLimit(limit) => redirect::Policy::limited(*limit as usize),
254                RedirectPolicy::FollowAll => redirect::Policy::limited(100),
255            });
256        }
257        let request = request.body(match body.0 {
258            http_client::Inner::Empty => reqwest::Body::default(),
259            http_client::Inner::Bytes(cursor) => cursor.into_inner().into(),
260            http_client::Inner::AsyncReader(stream) => {
261                reqwest::Body::wrap_stream(StreamReader::new(stream))
262            }
263        });
264
265        let handle = self.handle.clone();
266        async move {
267            let join_handle = handle.spawn(async { request.send().await });
268            let abort_handle = join_handle.abort_handle();
269            let _abort_on_drop = defer(move || abort_handle.abort());
270
271            let mut response = join_handle.await?.map_err(redact_error)?;
272
273            let headers = mem::take(response.headers_mut());
274            let mut builder = http::Response::builder()
275                .status(response.status().as_u16())
276                .version(response.version());
277            *builder.headers_mut().unwrap() = headers;
278
279            let bytes = response
280                .bytes_stream()
281                .map_err(futures::io::Error::other)
282                .into_async_read();
283            let body = http_client::AsyncBody::from_reader(bytes);
284
285            builder.body(body).map_err(|e| anyhow!(e))
286        }
287        .boxed()
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use std::io::{Read as _, Write as _};
294    use std::net::TcpListener;
295
296    use http_client::{AsyncBody, HttpClient, Method, Request as HttpRequest, Url};
297
298    use crate::ReqwestClient;
299
300    /// Regression test: `StreamReader::poll_next` used to drop the reader it
301    /// `take()`s whenever the reader returned `Poll::Pending`, so the next
302    /// poll reported end-of-stream and streamed request bodies were silently
303    /// truncated. Readers backed by real I/O (e.g. `async_fs::File`) return
304    /// `Pending` on their very first read, so their uploads sent zero bytes.
305    #[test]
306    fn test_streamed_body_survives_pending_reader() {
307        let payload: Vec<u8> = (0..30_000usize).map(|byte| (byte % 251) as u8).collect();
308
309        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
310        let address = listener.local_addr().unwrap();
311        let expected_payload = payload.clone();
312        let server = std::thread::spawn(move || {
313            let (mut stream, _) = listener.accept().unwrap();
314            let mut request = Vec::new();
315            let mut buffer = [0u8; 8192];
316            loop {
317                let read = stream.read(&mut buffer).unwrap();
318                assert_ne!(read, 0, "client closed the connection mid-request");
319                request.extend_from_slice(&buffer[..read]);
320                if let Some(position) = request.windows(4).position(|w| w == b"\r\n\r\n") {
321                    let body_start = position + 4;
322                    while request.len() - body_start < expected_payload.len() {
323                        let read = stream.read(&mut buffer).unwrap();
324                        assert_ne!(read, 0, "client closed the connection mid-body");
325                        request.extend_from_slice(&buffer[..read]);
326                    }
327                    assert_eq!(&request[body_start..], &expected_payload);
328                    break;
329                }
330            }
331            stream
332                .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\nconnection: close\r\n\r\n")
333                .unwrap();
334        });
335
336        // A reader that returns `Pending` before every chunk, like a reader
337        // backed by real I/O would.
338        struct PendingFirstReader {
339            data: std::io::Cursor<Vec<u8>>,
340            ready: bool,
341        }
342
343        impl futures::AsyncRead for PendingFirstReader {
344            fn poll_read(
345                mut self: std::pin::Pin<&mut Self>,
346                cx: &mut std::task::Context<'_>,
347                buf: &mut [u8],
348            ) -> std::task::Poll<std::io::Result<usize>> {
349                if self.ready {
350                    self.ready = false;
351                    std::task::Poll::Ready(self.data.read(buf))
352                } else {
353                    self.ready = true;
354                    cx.waker().wake_by_ref();
355                    std::task::Poll::Pending
356                }
357            }
358        }
359
360        let reader = PendingFirstReader {
361            data: std::io::Cursor::new(payload.clone()),
362            ready: false,
363        };
364
365        let client = ReqwestClient::new();
366        let request = HttpRequest::builder()
367            .method(Method::PUT)
368            .uri(format!("http://{address}/upload"))
369            .header("Content-Length", payload.len().to_string())
370            .body(AsyncBody::from_reader(reader))
371            .unwrap();
372        let response = futures::executor::block_on(client.send(request)).unwrap();
373        assert!(response.status().is_success());
374        server.join().unwrap();
375    }
376
377    #[test]
378    fn test_proxy_uri() {
379        let client = ReqwestClient::new();
380        assert_eq!(client.proxy(), None);
381
382        let proxy = Url::parse("http://localhost:10809").unwrap();
383        let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
384        assert_eq!(client.proxy(), Some(&proxy));
385
386        let proxy = Url::parse("https://localhost:10809").unwrap();
387        let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
388        assert_eq!(client.proxy(), Some(&proxy));
389
390        let proxy = Url::parse("socks4://localhost:10808").unwrap();
391        let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
392        assert_eq!(client.proxy(), Some(&proxy));
393
394        let proxy = Url::parse("socks4a://localhost:10808").unwrap();
395        let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
396        assert_eq!(client.proxy(), Some(&proxy));
397
398        let proxy = Url::parse("socks5://localhost:10808").unwrap();
399        let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
400        assert_eq!(client.proxy(), Some(&proxy));
401
402        let proxy = Url::parse("socks5h://localhost:10808").unwrap();
403        let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
404        assert_eq!(client.proxy(), Some(&proxy));
405    }
406
407    #[test]
408    fn test_invalid_proxy_uri() {
409        let proxy = Url::parse("socks://127.0.0.1:20170").unwrap();
410        let client = ReqwestClient::proxy_and_user_agent(Some(proxy), "test").unwrap();
411        assert!(
412            client.proxy.is_none(),
413            "An invalid proxy URL should add no proxy to the client!"
414        )
415    }
416}
417