1use std::sync::OnceLock;
4use std::sync::atomic::{AtomicU64, Ordering};
5use std::time::Duration;
6
7use indexmap::IndexMap;
8use reqwest::header::HeaderMap;
9use serde_json::Value;
10use tokio_util::sync::CancellationToken;
11
12use crate::sse;
13use crate::stream::StreamOptions;
14
15pub const ABORTED_BEFORE_RESPONSE: &str = "Request aborted";
17pub const ABORTED_DURING_STREAM: &str = "Request was aborted";
19pub const ABORTED_READ: &str = "This operation was aborted";
22
23const DEFAULT_MAX_RETRY_DELAY_MS: u64 = 60_000;
24
25static SETTINGS_PROXY: OnceLock<String> = OnceLock::new();
26
27pub const DEFAULT_IDLE_TIMEOUT_MS: u64 = 300_000;
29
30static IDLE_TIMEOUT_MS: AtomicU64 = AtomicU64::new(DEFAULT_IDLE_TIMEOUT_MS);
31
32pub fn set_idle_timeout_ms(ms: u64) {
36 IDLE_TIMEOUT_MS.store(ms, Ordering::Relaxed);
37}
38
39pub(crate) fn idle_timeout() -> Option<Duration> {
41 match IDLE_TIMEOUT_MS.load(Ordering::Relaxed) {
42 0 => None,
43 ms => Some(Duration::from_millis(ms)),
44 }
45}
46
47async fn within_idle_timeout<T>(future: impl std::future::Future<Output = T>) -> Option<T> {
50 match idle_timeout() {
51 Some(limit) => tokio::time::timeout(limit, future).await.ok(),
52 None => Some(future.await),
53 }
54}
55
56pub fn set_settings_proxy(proxy: Option<&str>) {
60 if let Some(proxy) = proxy.map(str::trim).filter(|proxy| !proxy.is_empty()) {
61 let _ = SETTINGS_PROXY.set(proxy.to_owned());
62 }
63}
64
65pub fn proxy_env() -> Vec<(&'static str, String)> {
68 let Some(proxy) = SETTINGS_PROXY.get() else {
69 return Vec::new();
70 };
71 ["HTTP_PROXY", "HTTPS_PROXY"]
72 .into_iter()
73 .filter(|name| std::env::var_os(name).is_none())
74 .map(|name| (name, proxy.clone()))
75 .collect()
76}
77
78pub fn client() -> &'static reqwest::Client {
81 static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
82 CLIENT.get_or_init(|| builder().build().unwrap_or_default())
83}
84
85pub(crate) fn upgrade_client() -> &'static reqwest::Client {
87 static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
88 CLIENT.get_or_init(|| builder().http1_only().build().unwrap_or_default())
89}
90
91fn builder() -> reqwest::ClientBuilder {
92 let mut builder = reqwest::Client::builder().user_agent(user_agent());
93 for (name, proxy) in proxy_env() {
94 let proxy = match name {
95 "HTTP_PROXY" => reqwest::Proxy::http(&proxy),
96 _ => reqwest::Proxy::https(&proxy),
97 };
98 if let Ok(proxy) = proxy {
99 builder = builder.proxy(proxy.no_proxy(reqwest::NoProxy::from_env()));
100 }
101 }
102 builder
103}
104
105pub fn user_agent() -> String {
107 format!(
108 "yapi/{} ({}; {})",
109 env!("CARGO_PKG_VERSION"),
110 std::env::consts::OS,
111 std::env::consts::ARCH
112 )
113}
114
115#[derive(Clone, Debug, Default)]
119pub(crate) struct Headers(pub(crate) Vec<(String, String)>);
120
121impl Headers {
122 pub(crate) fn set(&mut self, name: &str, value: Option<impl Into<String>>) {
124 self.0.retain(|(key, _)| !key.eq_ignore_ascii_case(name));
125 if let Some(value) = value {
126 self.0.push((name.to_owned(), value.into()));
127 }
128 }
129
130 pub(crate) fn extend(&mut self, headers: &IndexMap<String, Option<String>>) {
132 for (name, value) in headers {
133 self.set(name, value.as_deref());
134 }
135 }
136
137 pub(crate) fn extend_model(&mut self, headers: Option<&IndexMap<String, String>>) {
139 for (name, value) in headers.into_iter().flatten() {
140 self.set(name, Some(value.as_str()));
141 }
142 }
143
144 pub(crate) fn get(&self, name: &str) -> Option<&str> {
146 self.0
147 .iter()
148 .find(|(key, _)| key.eq_ignore_ascii_case(name))
149 .map(|(_, value)| value.as_str())
150 }
151
152 pub(crate) fn apply(&self, mut request: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
154 for (name, value) in &self.0 {
155 request = request.header(name.as_str(), value.as_str());
156 }
157 request
158 }
159}
160
161#[derive(Debug)]
163pub enum Failure {
164 Status {
166 status: u16,
168 body: String,
170 },
171 Connection(String),
173 Aborted,
175 RetryDelay {
177 delay: String,
179 status: u16,
181 body: String,
183 },
184}
185
186impl Failure {
187 pub fn plain_message(&self, sdk_message: impl FnOnce(u16, &str) -> String) -> Option<String> {
191 match self {
192 Failure::Status { .. } => None,
193 Failure::Connection(message) => Some(message.clone()),
194 Failure::Aborted => Some(ABORTED_BEFORE_RESPONSE.to_owned()),
195 Failure::RetryDelay {
196 delay,
197 status,
198 body,
199 } => Some(format!("{delay} {}", sdk_message(*status, body))),
200 }
201 }
202}
203
204fn retryable(failure: &Failure, headers: &HeaderMap) -> bool {
205 match failure {
206 Failure::Status { status, .. } => {
207 match headers.get("x-should-retry").and_then(|v| v.to_str().ok()) {
208 Some("true") => true,
209 Some("false") => false,
210 _ => matches!(status, 408 | 409 | 429) || *status >= 500,
211 }
212 }
213 Failure::Connection(_) => true,
214 Failure::Aborted | Failure::RetryDelay { .. } => false,
215 }
216}
217
218pub(crate) async fn post(
221 url: &str,
222 headers: &Headers,
223 body: &str,
224 timeout_ms: Option<u64>,
225 max_retries: u32,
226 max_retry_delay_ms: Option<u64>,
227 cancel: &CancellationToken,
228) -> Result<reqwest::Response, Failure> {
229 let options = StreamOptions {
230 max_retries,
231 max_retry_delay_ms,
232 cancel: cancel.clone(),
233 ..StreamOptions::default()
234 };
235 let build = || {
236 let request = headers.apply(client().post(url).body(body.to_owned()));
237 match timeout_ms {
238 Some(timeout) => request.timeout(Duration::from_millis(timeout)),
239 None => request,
240 }
241 };
242 send(build, &options).await
243}
244
245pub async fn send(
248 build: impl Fn() -> reqwest::RequestBuilder,
249 options: &StreamOptions,
250) -> Result<reqwest::Response, Failure> {
251 send_with(build, options, true).await
252}
253
254pub(crate) async fn send_with(
258 build: impl Fn() -> reqwest::RequestBuilder,
259 options: &StreamOptions,
260 use_headers: bool,
261) -> Result<reqwest::Response, Failure> {
262 let mut retry = 0;
263 loop {
264 let result = options
265 .cancel
266 .run_until_cancelled(attempt(build(), options.timeout_ms))
267 .await
268 .unwrap_or_else(|| Err(Box::new((Failure::Aborted, HeaderMap::new()))));
269 let (failure, mut headers) = match result {
270 Ok(response) => return Ok(response),
271 Err(failure) => *failure,
272 };
273 if !use_headers {
274 headers.clear();
275 }
276 if options.cancel.is_cancelled() {
277 return Err(Failure::Aborted);
278 }
279 if retry >= options.max_retries || !retryable(&failure, &headers) {
280 return Err(failure);
281 }
282 let delay = match retry_delay(&headers, retry, options.max_retry_delay_ms) {
283 Ok(delay) => delay,
284 Err(delay) => {
285 return Err(match failure {
286 Failure::Status { status, body } => Failure::RetryDelay {
287 delay,
288 status,
289 body,
290 },
291 other => other,
293 });
294 }
295 };
296 retry += 1;
297 options
298 .cancel
299 .run_until_cancelled(tokio::time::sleep(delay))
300 .await
301 .ok_or(Failure::Aborted)?;
302 }
303}
304
305type Attempt = Box<(Failure, HeaderMap)>;
307
308async fn attempt(
309 request: reqwest::RequestBuilder,
310 timeout_ms: Option<f64>,
311) -> Result<reqwest::Response, Attempt> {
312 let failed = |message: String| Box::new((Failure::Connection(message), HeaderMap::new()));
313 let send = within_idle_timeout(request.send());
317 let sdk_limit = timeout_ms
318 .filter(|ms| ms.is_finite() && *ms > 0.0)
319 .map(|ms| Duration::from_millis(ms as u64))
320 .filter(|limit| idle_timeout().is_none_or(|idle| *limit < idle));
321 let sent = match sdk_limit {
322 Some(limit) => tokio::time::timeout(limit, send)
323 .await
324 .map_err(|_| failed("Request timed out.".to_owned()))?,
325 None => send.await,
326 };
327 let response = match sent {
328 Some(Ok(response)) => response,
329 Some(Err(err)) => return Err(failed(connection_message(&err))),
330 None => return Err(failed("Connection error.".to_owned())),
331 };
332 let status = response.status().as_u16();
333 if (200..300).contains(&status) {
334 return Ok(response);
335 }
336 let headers = response.headers().clone();
337 let body = response.text().await.unwrap_or_default();
338 Err(Box::new((Failure::Status { status, body }, headers)))
339}
340
341fn connection_message(err: &reqwest::Error) -> String {
342 if err.is_timeout() {
343 "Request timed out.".to_owned()
344 } else {
345 "Connection error.".to_owned()
346 }
347}
348
349fn retry_delay(headers: &HeaderMap, retry: u32, max_ms: Option<u64>) -> Result<Duration, String> {
351 let header = |name: &str| headers.get(name).and_then(|value| value.to_str().ok());
352 let requested = header("retry-after-ms")
353 .and_then(|value| value.trim().parse::<f64>().ok())
354 .or_else(|| {
355 header("retry-after").and_then(|value| match value.trim().parse::<f64>() {
356 Ok(seconds) => Some(seconds * 1000.0),
357 Err(_) => crate::model_catalog::http_date_ms(value)
358 .map(|date| date as f64 - yapi_types::time::now_ms() as f64),
359 })
360 });
361 if let Some(ms) = requested.filter(|ms| ms.is_finite()) {
362 let max = max_ms.unwrap_or(DEFAULT_MAX_RETRY_DELAY_MS);
363 if max > 0 && ms > max as f64 {
364 return Err(format!(
365 "Server requested {}s retry delay (max: {}s).",
366 (ms / 1000.0).ceil(),
367 max.div_ceil(1000)
368 ));
369 }
370 return Ok(Duration::from_millis(ms.max(0.0) as u64));
371 }
372 let exponential = (0.5 * 2f64.powi(retry as i32)).min(8.0) * 1000.0;
373 Ok(Duration::from_millis((exponential * 0.875) as u64))
375}
376
377pub fn sdk_status_message(status: u16, error: Option<&Value>, raw: Option<&str>) -> String {
380 let message = match error {
381 Some(error) => match error.get("message") {
382 Some(Value::String(message)) if !message.is_empty() => Some(message.clone()),
383 Some(message) if is_truthy(message) => Some(yapi_types::json::stringify(message)),
384 _ if is_truthy(error) => Some(yapi_types::json::stringify(error)),
385 _ => raw.map(str::to_owned),
386 },
387 None => raw.map(str::to_owned),
388 };
389 match message.filter(|message| !message.is_empty()) {
390 Some(message) => format!("{status} {message}"),
391 None => format!("{status} status code (no body)"),
392 }
393}
394
395pub(crate) fn openai_status_message(status: u16, body: &str, prefix: Option<&str>) -> String {
398 let json = serde_json::from_str::<Value>(body).ok();
399 let error = json.as_ref().and_then(|json| json.get("error"));
400 provider_error_message(
401 &openai_sdk_message(status, body),
402 Some(status),
403 error,
404 prefix,
405 )
406}
407
408pub(crate) fn openai_sdk_message(status: u16, body: &str) -> String {
411 let json = serde_json::from_str::<Value>(body).ok();
412 let error = json.as_ref().and_then(|json| json.get("error"));
413 sdk_status_message(status, error, json.is_none().then_some(body))
414}
415
416pub fn sdk_error_message(error: &Value) -> String {
419 match error.get("message") {
420 Some(Value::String(message)) if !message.is_empty() => message.clone(),
421 Some(message) if is_truthy(message) => yapi_types::json::stringify(message),
422 _ if is_truthy(error) => yapi_types::json::stringify(error),
423 _ => "(no status code or body)".to_owned(),
424 }
425}
426
427pub const MALFORMED_SSE_JSON: &str = "Error reading response: malformed server-sent event JSON.";
429
430pub async fn next_openai_chunk(
434 reader: &mut SseReader,
435 cancel: &CancellationToken,
436) -> Result<Option<Value>, String> {
437 loop {
438 let Some(sse) = reader.next(cancel).await? else {
439 return Ok(None);
440 };
441 if sse.data == "[DONE]" {
442 return Ok(None);
443 }
444 let data: Value =
445 serde_json::from_str(&sse.data).map_err(|_| MALFORMED_SSE_JSON.to_owned())?;
446 match sse.event.as_deref() {
447 Some(event) if event.starts_with("thread.") => continue,
448 Some("error") => {
449 let error = data.get("error").filter(|error| !error.is_null());
450 return Err(sdk_error_message(error.unwrap_or(&data)));
451 }
452 _ => {}
453 }
454 if let Some(error) = data.get("error").filter(|error| is_truthy(error)) {
455 return Err(sdk_error_message(error));
456 }
457 return Ok(Some(data));
458 }
459}
460
461pub fn is_truthy(value: &Value) -> bool {
463 match value {
464 Value::Null => false,
465 Value::Bool(value) => *value,
466 Value::Number(number) => number.as_f64().is_some_and(|n| n != 0.0),
467 Value::String(text) => !text.is_empty(),
468 Value::Array(_) | Value::Object(_) => true,
469 }
470}
471
472const MAX_ERROR_BODY_CHARS: usize = 4000;
473
474pub fn provider_error_message(
479 message: &str,
480 status: Option<u16>,
481 body: Option<&Value>,
482 prefix: Option<&str>,
483) -> String {
484 let body = body
485 .filter(|body| body.as_object().is_some_and(|object| !object.is_empty()))
486 .map(yapi_types::json::stringify)
487 .map(|text| truncate_chars(text.trim(), MAX_ERROR_BODY_CHARS))
488 .filter(|body| !body.is_empty() && !message.contains(body.as_str()));
489 match (status, body, prefix) {
490 (Some(status), Some(body), Some(prefix)) => format!("{prefix} ({status}): {body}"),
491 (Some(status), Some(body), None) => format!("{status}: {body}"),
492 (Some(status), None, Some(prefix)) => format!("{prefix} ({status}): {message}"),
493 _ => message.to_owned(),
494 }
495}
496
497pub(crate) fn truncate_chars(text: &str, max: usize) -> String {
500 let len = yapi_types::js::len(text);
501 if len <= max {
502 return text.to_owned();
503 }
504 format!(
505 "{}... [truncated {} chars]",
506 yapi_types::js::slice(text, 0, max),
507 len - max
508 )
509}
510
511pub async fn read_chunk(
515 response: &mut reqwest::Response,
516 cancel: &CancellationToken,
517 interrupted: &str,
518) -> Result<Option<Vec<u8>>, String> {
519 if cancel.is_cancelled() {
520 return Err(ABORTED_DURING_STREAM.to_owned());
521 }
522 let Some(chunk) = cancel
525 .run_until_cancelled(within_idle_timeout(response.chunk()))
526 .await
527 else {
528 return Err(interrupted.to_owned());
529 };
530 chunk
531 .and_then(Result::ok)
532 .map(|chunk| chunk.map(|bytes| bytes.to_vec()))
533 .ok_or_else(|| "terminated".to_owned())
534}
535
536pub struct SseReader {
538 response: reqwest::Response,
539 decoder: sse::Decoder,
540 queue: std::collections::VecDeque<sse::Event>,
541 done: bool,
542 interrupted: &'static str,
543}
544
545impl SseReader {
546 pub fn new(response: reqwest::Response) -> SseReader {
549 SseReader {
550 response,
551 decoder: sse::Decoder::default(),
552 queue: Default::default(),
553 done: false,
554 interrupted: ABORTED_DURING_STREAM,
555 }
556 }
557
558 pub fn fetch(response: reqwest::Response) -> SseReader {
561 SseReader {
562 interrupted: ABORTED_READ,
563 ..SseReader::new(response)
564 }
565 }
566
567 pub async fn next(&mut self, cancel: &CancellationToken) -> Result<Option<sse::Event>, String> {
569 loop {
570 if let Some(event) = self.queue.pop_front() {
571 return Ok(Some(event));
572 }
573 if self.done {
574 return Ok(None);
575 }
576 match read_chunk(&mut self.response, cancel, self.interrupted).await? {
577 Some(bytes) => self.queue.extend(self.decoder.push(&bytes)),
578 None => {
579 self.done = true;
580 self.queue.extend(self.decoder.finish());
581 }
582 }
583 }
584 }
585}
586
587#[cfg(test)]
588mod tests {
589 use super::*;
590 use serde_json::json;
591
592 #[test]
593 fn proxy_fills_unset_variables() {
594 assert!(proxy_env().is_empty());
595 set_settings_proxy(Some(" "));
596 assert!(proxy_env().is_empty());
597 set_settings_proxy(Some(" http://proxy.test:8080 "));
598 set_settings_proxy(Some("http://ignored.test"));
599 let expected: Vec<(&str, String)> = ["HTTP_PROXY", "HTTPS_PROXY"]
600 .into_iter()
601 .filter(|name| std::env::var_os(name).is_none())
602 .map(|name| (name, "http://proxy.test:8080".to_owned()))
603 .collect();
604 assert_eq!(proxy_env(), expected);
605 }
606
607 #[test]
608 fn formats_like_sdks() {
609 let anthropic =
610 json!({"type":"error","error":{"type":"invalid_request_error","message":"bad"}});
611 assert_eq!(
612 sdk_status_message(400, Some(&anthropic), None),
613 r#"400 {"type":"error","error":{"type":"invalid_request_error","message":"bad"}}"#
614 );
615 assert_eq!(
616 sdk_status_message(404, Some(&json!({"message": "Not found"})), None),
617 "404 Not found"
618 );
619 assert_eq!(
620 sdk_status_message(502, None, Some("Bad gateway")),
621 "502 Bad gateway"
622 );
623 assert_eq!(
624 sdk_status_message(500, None, Some("")),
625 "500 status code (no body)"
626 );
627 assert_eq!(
628 sdk_status_message(500, None, None),
629 "500 status code (no body)"
630 );
631 }
632
633 #[test]
634 fn provider_errors_carry_body() {
635 let error = json!({"message": "Invalid model", "type": "invalid_request_error"});
636 assert_eq!(
637 provider_error_message("400 Invalid model", Some(400), Some(&error), None),
638 r#"400: {"message":"Invalid model","type":"invalid_request_error"}"#
639 );
640 assert_eq!(
641 provider_error_message("Connection error.", None, None, None),
642 "Connection error."
643 );
644 assert_eq!(
645 provider_error_message(
646 "400 Invalid model",
647 Some(400),
648 Some(&error),
649 Some("OpenAI API error")
650 ),
651 r#"OpenAI API error (400): {"message":"Invalid model","type":"invalid_request_error"}"#
652 );
653 assert_eq!(
654 provider_error_message(
655 "502 status code (no body)",
656 Some(502),
657 None,
658 Some("xai API error")
659 ),
660 "xai API error (502): 502 status code (no body)"
661 );
662 }
663
664 #[test]
665 fn honors_retry_after() {
666 let mut headers = HeaderMap::new();
667 headers.insert("retry-after", "2".parse().unwrap());
668 let failure = Failure::Status {
669 status: 429,
670 body: String::new(),
671 };
672 assert!(retryable(&failure, &headers));
673 assert_eq!(
674 retry_delay(&headers, 0, None).unwrap(),
675 Duration::from_secs(2)
676 );
677 assert!(retry_delay(&headers, 0, Some(1000)).is_err());
678 let iso = yapi_types::time::iso(yapi_types::time::now_ms() + 120_000);
680 let months = [
681 "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
682 ];
683 let month = months[iso[5..7].parse::<usize>().unwrap() - 1];
684 let date = format!(
685 "Mon, {} {month} {} {} GMT",
686 &iso[8..10],
687 &iso[..4],
688 &iso[11..19]
689 );
690 headers.insert("retry-after", date.parse().unwrap());
691 assert!(retry_delay(&headers, 0, Some(0)).unwrap() > Duration::from_secs(110));
692 assert_eq!(
693 retry_delay(&headers, 0, None).unwrap_err(),
694 "Server requested 120s retry delay (max: 60s)."
695 );
696 }
697}