#![forbid(unsafe_code)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![doc = concat!(
"[observability guide]: https://github.com/sagikazarmark/allama/blob/v",
env!("CARGO_PKG_VERSION"),
"/docs/guides/observability/README.md#trace-context-propagation",
)]
#![cfg_attr(
not(feature = "reqwest"),
doc = "[`ReqwestTransport`]: crate\n[`ReqwestTransport::new`]: crate\n[`DEFAULT_CONNECT_TIMEOUT`]: crate\n[`ReqwestTransport::from_client`]: crate"
)]
#![warn(missing_docs)]
#![warn(missing_debug_implementations)]
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
pub use bytes::Bytes;
use futures_core::Stream;
use futures_util::StreamExt;
pub mod config;
mod credential;
mod error;
pub mod framing;
pub mod ndjson;
mod options;
mod replay;
#[cfg(feature = "reqwest")]
mod reqwest;
mod retry_after;
pub mod sse;
mod timeout;
pub use credential::{
Credential, CredentialKind, CredentialPlacement, CredentialSlot, is_credential_header,
};
pub use error::{
ErrorEnvelope, error_context, error_envelope, error_envelope_value, malformed_response,
status_error,
};
pub use options::{HttpRequestOptions, OPTIONS_NAMESPACE};
pub use replay::{ReplayResponse, ReplayTransport};
#[cfg(feature = "reqwest")]
pub use reqwest::{DEFAULT_CONNECT_TIMEOUT, ReqwestTransport};
pub use retry_after::parse_retry_after;
pub use timeout::StreamTimeouts;
pub type SendFuture<'a> =
Pin<Box<dyn Future<Output = Result<HttpResponse, TransportError>> + Send + 'a>>;
pub type Body = Pin<Box<dyn Stream<Item = Result<Bytes, TransportError>> + Send>>;
pub trait HttpTransport: Send + Sync + 'static {
fn send(&self, request: HttpRequest) -> SendFuture<'_>;
}
impl<T: HttpTransport + ?Sized> HttpTransport for Arc<T> {
fn send(&self, request: HttpRequest) -> SendFuture<'_> {
T::send(self, request)
}
}
impl<T: HttpTransport + ?Sized> HttpTransport for Box<T> {
fn send(&self, request: HttpRequest) -> SendFuture<'_> {
T::send(self, request)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum Method {
Get,
Post,
}
impl Method {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Get => "GET",
Self::Post => "POST",
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct HttpRequest {
pub method: Method,
pub url: String,
pub headers: Headers,
pub body: Option<Bytes>,
pub timeout: Option<Duration>,
}
impl HttpRequest {
#[must_use]
pub fn get(url: impl Into<String>) -> Self {
Self {
method: Method::Get,
url: url.into(),
headers: Headers::new(),
body: None,
timeout: None,
}
}
#[must_use]
pub fn post(url: impl Into<String>, body: impl Into<Bytes>) -> Self {
Self {
method: Method::Post,
url: url.into(),
headers: Headers::new(),
body: Some(body.into()),
timeout: None,
}
}
#[must_use]
pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.insert(name, value);
self
}
#[must_use]
pub fn timeout(mut self, timeout: Option<Duration>) -> Self {
self.timeout = timeout;
self
}
#[must_use]
pub fn body_text(&self) -> Option<&str> {
std::str::from_utf8(self.body.as_deref()?).ok()
}
#[must_use]
pub fn sensitive_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.insert_sensitive(name, value);
self
}
#[must_use]
pub fn with_headers(mut self, headers: &Headers) -> Self {
for (name, _) in headers.iter() {
self.headers.remove(name);
}
for header in &headers.0 {
self.headers.0.push(header.clone());
}
self
}
pub fn apply_options(
mut self,
options: &allama::RequestOptions,
slot: &CredentialSlot,
) -> allama::Result<Self> {
let transport = options.resolve::<HttpRequestOptions>()?;
let transport = transport.as_deref();
if let Some(transport) = transport {
refuse_credential_headers(transport, slot)?;
}
let credential = transport.and_then(|transport| transport.credential_for(&slot.scope));
if slot.required && credential.is_none() {
return Err(allama::Error::configuration(format!(
"no credential is addressed to credential scope `{}`, and the provider has no \
configured credential: attach one with `HttpRequestOptions::credential`",
slot.scope
))
.authored());
}
let Some(transport) = transport else {
return Ok(self);
};
for (name, _) in transport.headers.iter() {
self.headers.remove(name);
}
for (name, value) in transport.headers.iter() {
self.headers.insert(name, value);
}
if let Some(credential) = credential {
self.place(credential, slot)?;
}
for (name, secret) in transport.header_credentials_for(&slot.scope) {
self.headers.remove(name);
self.headers.insert_sensitive(name, secret);
}
Ok(self)
}
fn place(&mut self, credential: &Credential, slot: &CredentialSlot) -> allama::Result<()> {
let (name, value) = match (&slot.placement, credential.kind()) {
(CredentialPlacement::Custom, _) => return Ok(()),
(CredentialPlacement::Bearer, CredentialKind::ApiKey) => (
"authorization",
format!("Bearer {}", credential.expose_secret()),
),
(CredentialPlacement::Header(name), CredentialKind::ApiKey) => {
(name.as_ref(), credential.expose_secret().to_owned())
}
_ => return Err(credential.unsupported(&slot.scope)),
};
self.headers.remove(name);
for replaced in &slot.replaces {
self.headers.remove(replaced);
}
self.headers.insert_sensitive(name, value);
Ok(())
}
}
fn refuse_credential_headers(
transport: &HttpRequestOptions,
slot: &CredentialSlot,
) -> allama::Result<()> {
for (name, _) in transport.headers.iter() {
if is_credential_header(name) {
return Err(allama::Error::configuration(format!(
"the plain header `{name}` carries a credential: send it with \
`HttpRequestOptions::credential`, addressed to a credential scope"
))
.authored());
}
if slot.owns_header(name) {
return Err(allama::Error::configuration(format!(
"the plain header `{name}` is where credential scope `{}` puts its credential: \
send it with `HttpRequestOptions::credential`",
slot.scope
))
.authored());
}
}
Ok(())
}
#[derive(Clone, Default, PartialEq, Eq)]
pub struct Headers(Vec<Header>);
impl Headers {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn insert(&mut self, name: impl Into<String>, value: impl Into<String>) {
self.0.push(Header {
name: name.into(),
value: value.into(),
sensitive: false,
});
}
pub fn insert_sensitive(&mut self, name: impl Into<String>, value: impl Into<String>) {
self.0.push(Header {
name: name.into(),
value: value.into(),
sensitive: true,
});
}
pub fn remove(&mut self, name: &str) {
self.0
.retain(|header| !header.name.eq_ignore_ascii_case(name));
}
#[must_use]
pub fn get(&self, name: &str) -> Option<&str> {
self.0
.iter()
.find(|header| header.name.eq_ignore_ascii_case(name))
.map(|header| header.value.as_str())
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
self.0
.iter()
.map(|header| (header.name.as_str(), header.value.as_str()))
}
#[must_use]
pub fn retry_after(&self, name: &str) -> Option<Duration> {
parse_retry_after(self.get(name), self.get("retry-after-ms"), now())
}
}
impl fmt::Debug for Headers {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_map()
.entries(self.0.iter().map(|header| (&header.name, header)))
.finish()
}
}
#[derive(Clone, PartialEq, Eq)]
struct Header {
name: String,
value: String,
sensitive: bool,
}
impl Header {
fn is_sensitive(&self) -> bool {
self.sensitive || is_sensitive_header(&self.name)
}
}
impl fmt::Debug for Header {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.is_sensitive() {
f.write_str("<redacted>")
} else {
fmt::Debug::fmt(&self.value, f)
}
}
}
#[must_use]
pub fn is_sensitive_header(name: &str) -> bool {
const FRAGMENTS: [&str; 4] = ["key", "token", "secret", "auth"];
let name = name.to_ascii_lowercase();
is_credential_header(&name)
|| name == "set-cookie"
|| FRAGMENTS.iter().any(|part| name.contains(part))
}
fn now() -> SystemTime {
let since_epoch = web_time::SystemTime::now()
.duration_since(web_time::UNIX_EPOCH)
.unwrap_or_default();
UNIX_EPOCH + since_epoch
}
impl<K: Into<String>, V: Into<String>> FromIterator<(K, V)> for Headers {
fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
let mut headers = Self::new();
for (name, value) in iter {
headers.insert(name, value);
}
headers
}
}
#[non_exhaustive]
pub struct HttpResponse {
pub status: u16,
pub headers: Headers,
body: Body,
}
impl fmt::Debug for HttpResponse {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("HttpResponse")
.field("status", &self.status)
.field("headers", &self.headers)
.finish_non_exhaustive()
}
}
impl HttpResponse {
#[must_use]
pub fn new(status: u16, headers: Headers, body: Body) -> Self {
Self {
status,
headers,
body,
}
}
#[must_use]
pub fn from_bytes(status: u16, headers: Headers, body: impl Into<Bytes>) -> Self {
let body = body.into();
Self::new(
status,
headers,
Box::pin(futures_util::stream::once(async move { Ok(body) })),
)
}
#[must_use]
pub fn is_success(&self) -> bool {
(200..300).contains(&self.status)
}
pub async fn bytes(self) -> Result<Bytes, TransportError> {
self.bytes_limited(usize::MAX).await
}
pub async fn text(self) -> Result<String, TransportError> {
self.text_limited(usize::MAX).await
}
pub async fn bytes_limited(self, limit: usize) -> Result<Bytes, TransportError> {
let mut body = self.body;
let mut buffer = Vec::new();
while let Some(chunk) = body.next().await {
let chunk = chunk?;
if chunk.len() > limit.saturating_sub(buffer.len()) {
return Err(TransportError::new(
TransportErrorKind::TooLarge,
format!("the response body is larger than the limit of {limit} bytes"),
)
.authored());
}
if buffer.is_empty() && body_is_done(&body) {
return Ok(chunk);
}
buffer.extend_from_slice(&chunk);
}
Ok(buffer.into())
}
pub async fn text_limited(self, limit: usize) -> Result<String, TransportError> {
let bytes = self.bytes_limited(limit).await?;
Ok(String::from_utf8_lossy(&bytes).into_owned())
}
pub async fn map_error(
self,
map: impl FnOnce(u16, &Headers, &str) -> allama::Error,
) -> allama::Error {
if (300..400).contains(&self.status) {
return redirect_error(self.status, &self.headers);
}
let (body, failure) =
read_error_body(self.body, MAX_ERROR_BODY_SIZE, ERROR_BODY_TIMEOUT).await;
let error = map(self.status, &self.headers, &String::from_utf8_lossy(&body));
let Some(failure) = failure else {
return error;
};
let note = format!("reading the error body failed: {failure}");
let message = match error.upstream_message() {
Some(message) => format!("{message} ({note})"),
None => note,
};
error.with_upstream_message(message)
}
#[must_use]
pub fn into_body(self) -> Body {
self.body
}
}
pub const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
pub const DEFAULT_MAX_RESPONSE_SIZE: usize = 256 * 1024 * 1024;
pub const MAX_ERROR_BODY_SIZE: usize = 64 * 1024;
pub const ERROR_BODY_TIMEOUT: Duration = Duration::from_secs(30);
pub const DEFAULT_USER_AGENT: &str = concat!("allama/", env!("CARGO_PKG_VERSION"));
#[cfg_attr(not(feature = "reqwest"), doc = "[`ReqwestTransport::new`]: crate")]
#[must_use]
pub fn redirect_error(status: u16, headers: &Headers) -> allama::Error {
let target = match headers.get("location").map(config::url_host) {
Some(Some(host)) => format!("to `{host}`"),
Some(None) => "within the same host".to_owned(),
None => "without a `Location`".to_owned(),
};
allama::Error::unexpected_response(format!(
"HTTP {status}: the server redirected the request {target}; redirects are not followed"
))
.with_context(error_context(status, headers))
.authored()
}
fn body_is_done(body: &Body) -> bool {
body.size_hint() == (0, Some(0))
}
async fn read_error_body(
mut body: Body,
limit: usize,
timeout: Duration,
) -> (Vec<u8>, Option<TransportError>) {
let mut buffer = Vec::new();
let failure = {
let read = std::pin::pin!(async {
while let Some(chunk) = body.next().await {
let chunk = match chunk {
Ok(chunk) => chunk,
Err(error) => return Some(error),
};
let room = limit - buffer.len();
buffer.extend_from_slice(&chunk[..chunk.len().min(room)]);
if buffer.len() == limit {
break;
}
}
None
});
let deadline = futures_timer::Delay::new(timeout);
match futures_util::future::select(read, deadline).await {
futures_util::future::Either::Left((failure, _)) => failure,
futures_util::future::Either::Right(((), _)) => Some(
TransportError::new(
TransportErrorKind::Timeout,
format!("the error body did not end within {timeout:?}"),
)
.authored(),
),
}
};
(buffer, failure)
}
#[derive(Debug)]
pub struct TransportError {
kind: TransportErrorKind,
message: String,
authored: bool,
source: Option<Box<dyn std::error::Error + Send + Sync>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum TransportErrorKind {
Timeout,
InvalidRequest,
Connect,
Body,
TooLarge,
Other,
}
impl TransportError {
#[must_use]
pub fn new(kind: TransportErrorKind, message: impl Into<String>) -> Self {
Self {
kind,
message: message.into(),
authored: false,
source: None,
}
}
#[must_use]
pub fn authored(mut self) -> Self {
self.authored = true;
self
}
#[must_use]
pub fn with_source(mut self, source: impl std::error::Error + Send + Sync + 'static) -> Self {
self.source = Some(Box::new(source));
self
}
#[must_use]
pub fn kind(&self) -> TransportErrorKind {
self.kind
}
#[must_use]
pub fn into_error(mut self) -> allama::Error {
let source = self.source.take();
if self.authored {
let error = match self.kind {
TransportErrorKind::Timeout => allama::Error::timeout(self.message),
TransportErrorKind::InvalidRequest => allama::Error::invalid_request(self.message),
TransportErrorKind::TooLarge => allama::Error::unexpected_response(self.message),
_ => allama::Error::transport(self.message),
};
return with_source(error.authored(), source);
}
let error = match self.kind {
TransportErrorKind::Timeout => allama::Error::timeout("the request timed out"),
TransportErrorKind::InvalidRequest => {
allama::Error::invalid_request("the request could not be built")
}
TransportErrorKind::Connect => allama::Error::transport("the request got no response"),
TransportErrorKind::Body => {
allama::Error::transport("reading the response body failed")
}
TransportErrorKind::TooLarge => {
allama::Error::unexpected_response("the response body is larger than the limit")
}
TransportErrorKind::Other => allama::Error::transport("the transport failed"),
};
with_source(error.authored().with_upstream_message(self.message), source)
}
#[must_use]
pub fn into_stream_error(mut self) -> allama::Error {
let source = self.source.take();
if self.authored {
let error = match self.kind {
TransportErrorKind::Timeout => allama::Error::timeout(self.message),
TransportErrorKind::TooLarge => allama::Error::stream_protocol(self.message),
_ => allama::Error::stream(allama::StreamErrorKind::Transport, self.message),
};
return with_source(error.authored(), source);
}
let error = match self.kind {
TransportErrorKind::Timeout => allama::Error::timeout("the stream timed out"),
TransportErrorKind::TooLarge => {
allama::Error::stream_protocol("the stream is larger than the limit")
}
_ => allama::Error::stream(
allama::StreamErrorKind::Transport,
"reading the stream failed",
),
};
with_source(error.authored().with_upstream_message(self.message), source)
}
}
fn with_source(
error: allama::Error,
source: Option<Box<dyn std::error::Error + Send + Sync>>,
) -> allama::Error {
match source {
Some(source) => error.with_source(source),
None => error,
}
}
impl fmt::Display for TransportError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.kind {
TransportErrorKind::Timeout => write!(f, "request timed out: {}", self.message),
_ => f.write_str(&self.message),
}
}
}
impl std::error::Error for TransportError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
self.source
.as_deref()
.map(|e| e as &(dyn std::error::Error + 'static))
}
}
impl From<TransportError> for allama::Error {
fn from(err: TransportError) -> Self {
err.into_error()
}
}
#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
use super::*;
#[test]
fn headers_are_case_insensitive() {
let headers: Headers = [("Request-Id", "req_1"), ("retry-after", "2")]
.into_iter()
.collect();
assert_eq!(headers.get("request-id"), Some("req_1"));
assert_eq!(headers.get("RETRY-AFTER"), Some("2"));
assert_eq!(headers.get("missing"), None);
}
#[test]
fn debug_output_redacts_credentials() {
let request = HttpRequest::post("https://x", "{}")
.header("Authorization", "Bearer sk-live-SECRET1")
.header("x-api-key", "SECRET2")
.header("x-goog-api-key", "SECRET3")
.header("Proxy-Authorization", "Basic SECRET4")
.header("cookie", "session=SECRET5")
.header("x-gateway-token", "SECRET6")
.header("X-Client-Secret", "SECRET7")
.header("api-key", "SECRET8")
.header("x-request-id", "req_1");
let debug = format!("{request:?}");
assert!(!debug.contains("SECRET"), "{debug}");
assert!(debug.contains(r#""x-request-id": "req_1""#), "{debug}");
assert!(debug.contains(r#""Authorization": <redacted>"#), "{debug}");
let response = HttpResponse::from_bytes(
200,
[
("set-cookie", "session=SECRET9"),
("content-type", "text/plain"),
]
.into_iter()
.collect(),
"",
);
let debug = format!("{response:?}");
assert!(!debug.contains("SECRET"), "{debug}");
assert!(debug.contains("text/plain"), "{debug}");
let signed =
HttpRequest::post("https://x", "{}").sensitive_header("x-signature", "SECRET10");
assert_eq!(signed.headers.get("x-signature"), Some("SECRET10"));
assert!(!format!("{signed:?}").contains("SECRET"));
}
#[test]
fn sensitive_headers_are_matched_by_name_and_fragment() {
for name in [
"AUTHORIZATION",
"cookie",
"Set-Cookie",
"x-auth",
"openai-api-key",
] {
assert!(is_sensitive_header(name), "{name}");
}
for name in ["content-type", "x-request-id", "retry-after", "x-tenant-id"] {
assert!(!is_sensitive_header(name), "{name}");
}
}
#[test]
fn retry_after_reads_the_named_header_or_milliseconds() {
let headers: Headers = [("x-retry", "3")].into_iter().collect();
assert_eq!(headers.retry_after("x-retry"), Some(Duration::from_secs(3)));
assert_eq!(headers.retry_after("retry-after"), None);
let headers: Headers = [("retry-after-ms", "1500")].into_iter().collect();
assert_eq!(
headers.retry_after("retry-after"),
Some(Duration::from_millis(1500))
);
let past: Headers = [("retry-after", "Sun, 06 Nov 1994 08:49:37 GMT")]
.into_iter()
.collect();
assert_eq!(past.retry_after("retry-after"), Some(Duration::ZERO));
}
#[test]
fn requests_build_fluently() {
let request = HttpRequest::post("https://x", "{}")
.header("authorization", "Bearer k")
.timeout(Some(Duration::from_secs(5)));
assert_eq!(request.method.as_str(), "POST");
assert_eq!(request.headers.get("Authorization"), Some("Bearer k"));
assert_eq!(request.body_text(), Some("{}"));
assert!(HttpRequest::get("https://x").body.is_none());
}
fn bearer_slot() -> CredentialSlot {
CredentialSlot::new("openai", CredentialPlacement::Bearer)
}
fn options(transport: HttpRequestOptions) -> allama::RequestOptions {
let mut options = allama::RequestOptions::new();
options.insert(transport);
options
}
fn values<'a>(request: &'a HttpRequest, name: &str) -> Vec<&'a str> {
request
.headers
.iter()
.filter(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| value)
.collect()
}
#[test]
fn request_options_replace_same_named_headers_and_add_others() {
let base = || {
HttpRequest::post("https://x", "{}")
.header("Authorization", "Bearer provider")
.header("x-tag", "one")
.header("x-tag", "two")
.header("content-type", "application/json")
};
let untouched = base()
.apply_options(&allama::RequestOptions::new(), &bearer_slot())
.unwrap();
assert_eq!(untouched, base());
let applied = base()
.apply_options(
&options(
HttpRequestOptions::new()
.header("X-TAG", "three")
.header("x-tenant-id", "acme")
.header("Idempotency-Key", "idem-1"),
),
&bearer_slot(),
)
.unwrap();
let headers: Vec<_> = applied.headers.iter().collect();
assert_eq!(
headers,
[
("Authorization", "Bearer provider"),
("content-type", "application/json"),
("X-TAG", "three"),
("x-tenant-id", "acme"),
("Idempotency-Key", "idem-1"),
]
);
let mut invalid = allama::RequestOptions::new();
invalid.insert_json(
OPTIONS_NAMESPACE,
serde_json::json!({"headers": {"x-api-key": "k"}}),
);
assert!(matches!(
base().apply_options(&invalid, &bearer_slot()).unwrap_err(),
allama::Error::InvalidRequest { .. }
));
}
#[test]
fn the_scope_credential_replaces_the_configured_one_in_its_slot() {
let transport = HttpRequestOptions::new()
.credential("openai", Credential::api_key("TENANT-SECRET"))
.credential("anthropic", Credential::api_key("FOREIGN-SECRET"));
let bearer = HttpRequest::post("https://x", "{}")
.header("authorization", "Bearer provider")
.apply_options(&options(transport.clone()), &bearer_slot())
.unwrap();
assert_eq!(values(&bearer, "authorization"), ["Bearer TENANT-SECRET"]);
let azure = CredentialSlot::new("openai", CredentialPlacement::Header("api-key".into()))
.replaces("authorization");
let header = HttpRequest::post("https://x", "{}")
.header("authorization", "Bearer entra-token")
.header("API-KEY", "configured")
.apply_options(&options(transport.clone()), &azure)
.unwrap();
assert_eq!(values(&header, "api-key"), ["TENANT-SECRET"]);
assert!(values(&header, "authorization").is_empty());
for request in [&bearer, &header] {
let debug = format!("{request:?}");
assert!(!debug.contains("SECRET"), "{debug}");
assert!(
!request.headers.iter().any(|(_, v)| v.contains("FOREIGN")),
"{request:?}"
);
}
let custom = CredentialSlot::new("openai", CredentialPlacement::Custom);
let signed = HttpRequest::post("https://x", "{}")
.header("authorization", "configured")
.apply_options(&options(transport.clone()), &custom)
.unwrap();
assert_eq!(values(&signed, "authorization"), ["configured"]);
assert_eq!(
transport
.credential_for("openai")
.map(Credential::expose_secret),
Some("TENANT-SECRET")
);
}
#[test]
fn header_credentials_reach_only_their_scope_and_replace_same_named_headers() {
let transport = HttpRequestOptions::new()
.credential(
"openai",
Credential::header("cf-aig-authorization", "Bearer GATEWAY-SECRET"),
)
.credential(
"anthropic",
Credential::header("x-other-gw", "OTHER-SECRET"),
);
let applied = HttpRequest::post("https://x", "{}")
.header("authorization", "Bearer provider")
.header("cf-aig-authorization", "Bearer configured-gateway")
.apply_options(&options(transport), &bearer_slot())
.unwrap();
assert_eq!(
values(&applied, "cf-aig-authorization"),
["Bearer GATEWAY-SECRET"]
);
assert_eq!(values(&applied, "authorization"), ["Bearer provider"]);
assert!(values(&applied, "x-other-gw").is_empty());
let debug = format!("{applied:?}");
assert!(!debug.contains("SECRET"), "{debug}");
let custom_name = HttpRequest::post("https://x", "{}")
.apply_options(
&options(
HttpRequestOptions::new()
.credential("openai", Credential::header("x-route", "ROUTE-SECRET")),
),
&bearer_slot(),
)
.unwrap();
assert_eq!(values(&custom_name, "x-route"), ["ROUTE-SECRET"]);
assert!(!format!("{custom_name:?}").contains("ROUTE-SECRET"));
}
#[test]
fn plain_credential_headers_are_refused_as_configuration_errors() {
let nimbus = CredentialSlot::new(
"nimbus",
CredentialPlacement::Header("x-nimbus-session".into()),
)
.replaces("x-nimbus-legacy");
for (name, slot) in [
("Authorization", &nimbus),
("x-api-key", &nimbus),
("cf-aig-authorization", &nimbus),
("ocp-apim-subscription-key", &nimbus),
("x-amz-security-token", &nimbus),
("cookie", &nimbus),
("X-Nimbus-Session", &nimbus),
("x-nimbus-legacy", &nimbus),
] {
let err = HttpRequest::post("https://x", "{}")
.apply_options(
&options(HttpRequestOptions::new().header(name, "PLAIN-SECRET")),
slot,
)
.unwrap_err();
assert!(
matches!(err, allama::Error::Configuration { .. }),
"{err:?}"
);
assert!(!err.is_retryable(), "{name}");
let text = format!("{err} {err:?} {}", err.as_log());
assert!(!text.contains("PLAIN-SECRET"), "{text}");
assert!(text.contains(name), "{text}");
}
assert!(
HttpRequest::post("https://x", "{}")
.apply_options(
&options(HttpRequestOptions::new().header("x-nimbus-session", "v")),
&bearer_slot(),
)
.is_ok()
);
}
#[test]
fn a_required_slot_refuses_a_request_without_its_credential() {
let slot = bearer_slot().required(true);
for transport in [
None,
Some(HttpRequestOptions::new().credential("other", Credential::api_key("k"))),
Some(
HttpRequestOptions::new()
.credential("openai", Credential::header("cf-aig-authorization", "g")),
),
] {
let options = transport.map(options).unwrap_or_default();
let err = HttpRequest::post("https://x", "{}")
.apply_options(&options, &slot)
.unwrap_err();
assert!(
matches!(err, allama::Error::Configuration { .. }),
"{err:?}"
);
assert!(err.to_string().contains("`openai`"), "{err}");
}
let sent = HttpRequest::post("https://x", "{}")
.apply_options(
&options(HttpRequestOptions::new().credential("openai", Credential::api_key("k"))),
&slot,
)
.unwrap();
assert_eq!(sent.headers.get("authorization"), Some("Bearer k"));
}
#[test]
fn a_later_credential_replaces_the_same_kind_in_its_scope() {
let transport = HttpRequestOptions::new()
.credential("openai", Credential::api_key("first"))
.credential("openai", Credential::header("x-gw", "gw-1"))
.credential("OPENAI", Credential::api_key("other-scope"))
.credential("openai", Credential::api_key("second"))
.credential("openai", Credential::header("X-GW", "gw-2"));
assert_eq!(
transport
.credential_for("openai")
.map(Credential::expose_secret),
Some("second")
);
let scopes: Vec<_> = transport
.credentials()
.map(|(scope, credential)| (scope, credential.expose_secret()))
.collect();
assert_eq!(
scopes,
[
("OPENAI", "other-scope"),
("openai", "second"),
("openai", "gw-2")
]
);
}
#[tokio::test]
async fn bodies_are_read_whole_or_streamed() {
let response = HttpResponse::new(
200,
Headers::new(),
Box::pin(futures_util::stream::iter([
Ok(Bytes::from_static(b"he")),
Ok(Bytes::from_static(b"llo")),
])),
);
assert!(response.is_success());
assert_eq!(response.text().await.unwrap(), "hello");
let failing = HttpResponse::new(
200,
Headers::new(),
Box::pin(futures_util::stream::iter([
Ok(Bytes::from_static(b"he")),
Err(TransportError::new(TransportErrorKind::Body, "reset")),
])),
);
assert_eq!(
failing.bytes().await.unwrap_err().kind(),
TransportErrorKind::Body
);
}
fn chunks(parts: Vec<Result<Bytes, TransportError>>) -> Body {
Box::pin(futures_util::stream::iter(parts))
}
#[tokio::test]
async fn limited_reads_stop_before_passing_the_limit() {
let body = || {
chunks(vec![
Ok(Bytes::from_static(b"0123")),
Ok(Bytes::from_static(b"4567")),
Ok(Bytes::from_static(b"89")),
])
};
let read = |limit| HttpResponse::new(200, Headers::new(), body()).bytes_limited(limit);
assert_eq!(read(10).await.unwrap(), "0123456789");
assert_eq!(read(usize::MAX).await.unwrap(), "0123456789");
for limit in [0, 3, 4, 9] {
let err = read(limit).await.unwrap_err();
assert_eq!(err.kind(), TransportErrorKind::TooLarge, "{limit}");
let err = err.into_error();
assert!(
matches!(err, allama::Error::UnexpectedResponse { .. }),
"{err:?}"
);
assert!(!err.is_retryable());
assert!(
err.as_log()
.to_string()
.contains(&format!("limit of {limit} bytes")),
"{}",
err.as_log()
);
}
let text = HttpResponse::from_bytes(200, Headers::new(), "héllo")
.text_limited(6)
.await
.unwrap();
assert_eq!(text, "héllo");
assert!(matches!(
TransportError::new(TransportErrorKind::TooLarge, "x").into_stream_error(),
allama::Error::Stream {
kind: allama::StreamErrorKind::Protocol,
..
}
));
}
fn map_http(status: u16, _: &Headers, body: &str) -> allama::Error {
status_error(status, format!("HTTP {status}"))
.authored()
.with_body(body.to_owned())
}
#[tokio::test]
async fn error_bodies_are_read_up_to_the_cap() {
let endless: Body = Box::pin(futures_util::stream::repeat_with(|| {
Ok(Bytes::from(vec![b'x'; 16 * 1024]))
}));
let err = HttpResponse::new(503, Headers::new(), endless)
.map_error(map_http)
.await;
assert!(matches!(err, allama::Error::Overloaded { .. }), "{err:?}");
assert_eq!(err.upstream_message(), None);
let (body, failure) = read_error_body(
Box::pin(futures_util::stream::repeat_with(|| {
Ok(Bytes::from_static(b"abc"))
})),
MAX_ERROR_BODY_SIZE,
ERROR_BODY_TIMEOUT,
)
.await;
assert_eq!(body.len(), MAX_ERROR_BODY_SIZE);
assert!(failure.is_none());
}
#[tokio::test]
async fn error_body_failures_are_recorded() {
let stalled: Body = Box::pin(
futures_util::stream::iter([Ok(Bytes::from_static(b"{\"err"))])
.chain(futures_util::stream::pending()),
);
let (body, failure) =
read_error_body(stalled, MAX_ERROR_BODY_SIZE, Duration::from_millis(20)).await;
assert_eq!(body, b"{\"err");
assert_eq!(failure.unwrap().kind(), TransportErrorKind::Timeout);
let reset = HttpResponse::new(
500,
Headers::new(),
chunks(vec![
Ok(Bytes::from_static(b"partial")),
Err(TransportError::new(
TransportErrorKind::Body,
"connection reset",
)),
]),
);
let err = reset.map_error(map_http).await;
assert_eq!(
err.upstream_message(),
Some("reading the error body failed: connection reset")
);
assert_eq!(err.context().and_then(|c| c.capped_body()), Some("partial"));
let enveloped = HttpResponse::new(
500,
Headers::new(),
chunks(vec![Err(TransportError::new(
TransportErrorKind::Body,
"reset",
))]),
)
.map_error(|status, headers, body| {
map_http(status, headers, body).with_upstream_message("upstream")
})
.await;
assert_eq!(
enveloped.upstream_message(),
Some("upstream (reading the error body failed: reset)")
);
}
#[tokio::test]
async fn redirects_are_unexpected_responses_naming_the_host() {
let mut called = false;
let headers: Headers = [
(
"location",
"https://u:pw@Other.Example:8443/v1/x?token=SECRET",
),
("x-request-id", "req_1"),
]
.into_iter()
.collect();
let err = HttpResponse::from_bytes(308, headers, "")
.map_error(|status, headers, body| {
called = true;
map_http(status, headers, body)
})
.await;
assert!(!called);
assert!(
matches!(err, allama::Error::UnexpectedResponse { .. }),
"{err:?}"
);
assert_eq!(err.status(), Some(308));
assert_eq!(err.request_id(), Some("req_1"));
let log = err.as_log().to_string();
assert!(log.contains("`other.example`"), "{log}");
for text in [log, err.to_string(), format!("{err:?}")] {
assert!(!text.contains("SECRET") && !text.contains("pw"), "{text}");
}
let relative = redirect_error(301, &[("location", "/v2/chat")].into_iter().collect());
assert!(
relative.to_string().contains("within the same host"),
"{relative}"
);
let bare = redirect_error(302, &Headers::new());
assert!(bare.to_string().contains("without a `Location`"), "{bare}");
}
#[test]
fn with_headers_replaces_same_named_headers() {
let defaults: Headers = [("User-Agent", "a/1"), ("x-route", "eu")]
.into_iter()
.collect();
let request = HttpRequest::post("https://x", "{}")
.header("content-type", "application/json")
.header("user-agent", "old")
.sensitive_header("x-signed", "SECRET")
.with_headers(&defaults);
assert_eq!(
request.headers.iter().collect::<Vec<_>>(),
[
("content-type", "application/json"),
("x-signed", "SECRET"),
("User-Agent", "a/1"),
("x-route", "eu"),
]
);
assert!(!format!("{request:?}").contains("SECRET"));
}
#[test]
fn errors_map_to_portable_errors() {
let timeout = TransportError::new(TransportErrorKind::Timeout, "after 5s");
assert_eq!(timeout.to_string(), "request timed out: after 5s");
assert!(matches!(
allama::Error::from(timeout),
allama::Error::Timeout { .. }
));
assert!(matches!(
TransportError::new(TransportErrorKind::InvalidRequest, "bad url").into_error(),
allama::Error::InvalidRequest { .. }
));
assert!(matches!(
TransportError::new(TransportErrorKind::Connect, "refused").into_error(),
allama::Error::Transport { .. }
));
}
#[test]
fn transport_text_stays_out_of_the_log() {
for kind in [
TransportErrorKind::Timeout,
TransportErrorKind::InvalidRequest,
TransportErrorKind::Connect,
TransportErrorKind::Body,
TransportErrorKind::Other,
] {
let err = TransportError::new(kind, "https://x.example/v1?key=SECRET").into_error();
let line = err.as_log().to_string();
assert!(!line.contains("SECRET"), "{line}");
assert!(line.contains("message="), "{line}");
assert_eq!(
err.upstream_message(),
Some("https://x.example/v1?key=SECRET")
);
assert!(err.to_string().contains("SECRET"));
}
}
}
#[cfg(all(test, target_arch = "wasm32"))]
mod wasm_tests {
use wasm_bindgen_test::wasm_bindgen_test;
use super::*;
#[wasm_bindgen_test]
fn retry_after_reads_the_clock() {
let past: Headers = [("retry-after", "Sun, 06 Nov 1994 08:49:37 GMT")]
.into_iter()
.collect();
assert_eq!(past.retry_after("retry-after"), Some(Duration::ZERO));
let future: Headers = [("retry-after", "Fri, 31 Dec 9999 23:59:59 GMT")]
.into_iter()
.collect();
assert!(future.retry_after("retry-after").unwrap() > Duration::ZERO);
}
}