Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 105 additions & 4 deletions rig-core/src/http_client/mod.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
use crate::http_client::sse::BoxedStream;
use crate::http_client::{sse::BoxedStream, util::Escape};
use bytes::Bytes;
use http::StatusCode;
pub use http::{HeaderMap, HeaderValue, Method, Request, Response, Uri, request::Builder};
use http::{StatusCode, request::Parts};
use reqwest::{Body, multipart::Form};
use std::sync::atomic::{AtomicUsize, Ordering};
use tracing::Level;

pub mod retry;
pub mod sse;
pub mod util;

use std::pin::Pin;

Expand Down Expand Up @@ -75,6 +78,46 @@ pub fn with_bearer_auth(req: Builder, auth: &str) -> Result<Builder> {
Ok(req.header("Authorization", auth_header))
}

static LOG_HTTP_BODY_MAX: AtomicUsize = AtomicUsize::new(8 * 1024);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not really a fan of this static global. Is there a way we can do this per-type impl?


/// Extension trait for client builders to configure HTTP logging.
///
/// # Example
/// ```
/// use rig::prelude::*;
/// use rig::providers::openai;
///
/// let client = openai::Client::builder("api-key")
/// .max_log_body_preview(8192)
/// .build();
/// ```
pub trait HttpLogConfigExt {
/// Set the maximum number of bytes to preview from the body when logging in the `TRACE` level.
/// Defaults to 8192 bytes. Set to 0 to disable body preview logging.
///
/// This method can be called on any client builder to configure HTTP logging before building the client.
fn max_log_body_preview(self, max_preview_bytes: usize) -> Self;
}

impl<T> HttpLogConfigExt for T
where
T: Sized,
{
fn max_log_body_preview(self, max_preview_bytes: usize) -> Self {
LOG_HTTP_BODY_MAX.store(max_preview_bytes, Ordering::Relaxed);
self
}
}

/// Set the maximum number of bytes to preview from the body when logging in the `TRACE` level. Defaults to 8192 bytes. Set to 0 to disable body preview logging.
pub fn set_max_log_body_preview(max_preview_bytes: usize) {
LOG_HTTP_BODY_MAX.store(max_preview_bytes, Ordering::Relaxed);
}

fn body_preview_len() -> usize {
LOG_HTTP_BODY_MAX.load(Ordering::Relaxed)
}

/// A helper trait to make generic requests (both regular and SSE) possible.
pub trait HttpClientExt: WasmCompatSend + WasmCompatSync {
/// Send a HTTP request, get a response back (as bytes). Response must be able to be turned back into Bytes.
Expand Down Expand Up @@ -116,10 +159,14 @@ impl HttpClientExt for reqwest::Client {
U: From<Bytes> + WasmCompatSend,
{
let (parts, body) = req.into_parts();

let body_bytes: Bytes = body.into();
log_request(&parts, &body_bytes);

let req = self
.request(parts.method, parts.uri.to_string())
.headers(parts.headers)
.body(body.into());
.body(body_bytes);

async move {
let response = req.send().await.map_err(instance_error)?;
Expand Down Expand Up @@ -159,6 +206,9 @@ impl HttpClientExt for reqwest::Client {
U: WasmCompatSend + 'static,
{
let (parts, body) = req.into_parts();

log_headers(&parts);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to make an on/off toggle for this?


let req = self
.request(parts.method, parts.uri.to_string())
.headers(parts.headers)
Expand Down Expand Up @@ -202,10 +252,13 @@ impl HttpClientExt for reqwest::Client {
{
let (parts, body) = req.into_parts();

let body_bytes: Bytes = body.into();
log_request(&parts, &body_bytes);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to make an on/off toggle for this?


let req = self
.request(parts.method, parts.uri.to_string())
.headers(parts.headers)
.body(body.into())
.body(body_bytes)
.build()
.map_err(|x| Error::Instance(x.into()))
.unwrap();
Expand Down Expand Up @@ -246,3 +299,51 @@ impl HttpClientExt for reqwest::Client {
}
}
}

fn log_request(parts: &Parts, body: &Bytes) {
if tracing::enabled!(Level::TRACE) {
// Redact sensitive headers (e.g., Authorization) for logging
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can be removed.

It's pretty obvious what this is doing

let mut redacted_headers = parts.headers.clone();
redacted_headers.remove("Authorization");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You will also need to remove any headers that contain "-key" (as this is often used to hold API keys, etc).

let preview_len = body_preview_len();

if preview_len > 0 {
let shown = std::cmp::min(preview_len, body.len());
let preview = Escape::new(&body[..shown]);
tracing::trace!(
target: "rig::http",
method = %parts.method,
uri = %parts.uri,
body_len = body.len(),
body_preview_len = shown,
headers = ?redacted_headers,
body_preview = ?preview,
"sending HTTP request"
);
} else {
tracing::trace!(
target: "rig::http",
method = %parts.method,
uri = %parts.uri,
body_len = body.len(),
headers = ?redacted_headers,
"sending HTTP request"
);
}
}
}

fn log_headers(parts: &Parts) {
if tracing::enabled!(Level::TRACE) {
// Redact sensitive headers (e.g., Authorization) for logging
let mut redacted_headers = parts.headers.clone();
redacted_headers.remove("Authorization");
tracing::trace!(
target: "rig::http",
method = %parts.method,
uri = %parts.uri,
headers = ?redacted_headers,
"sending HTTP request"
);
}
}
41 changes: 41 additions & 0 deletions rig-core/src/http_client/util.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
use std::fmt;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please rename this file to escape.rs rather than util.rs.

Utility/helper files can generally lead to code smells if not well-maintained. I would like to be optimistic about this, but realistically it'll be a matter of when it will be a code smell and not if - so the best solution is to eliminate the problem entirely.

use std::str;

/// A helper struct to escape bytes for logging.
pub(crate) struct Escape<'a>(&'a [u8]);

impl<'a> Escape<'a> {
pub(crate) fn new(bytes: &'a [u8]) -> Self {
Escape(bytes)
}
}

impl fmt::Debug for Escape<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
// For valid UTF-8 strings, output directly for better readability
if let Ok(s) = str::from_utf8(self.0) {
return write!(f, "{}", s);
}
write!(f, "{}", self)
}
}

impl fmt::Display for Escape<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for &c in self.0 {
match c {
b'\n' => write!(f, "\\n")?,
b'\r' => write!(f, "\\r")?,
b'\t' => write!(f, "\\t")?,
b'\\' => write!(f, "\\\\")?,
b'"' => write!(f, "\\\"")?,
b'\0' => write!(f, "\\0")?,
// ASCII printable (0x20-0x7e, excluding space which is 0x20)
c if (0x20..0x7f).contains(&c) => write!(f, "{}", c as char)?,
// Non-printable bytes
c => write!(f, "\\x{c:02x}")?,
}
}
Ok(())
}
}
2 changes: 2 additions & 0 deletions rig-core/src/prelude.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,3 +16,5 @@ pub use crate::client::image_generation::ImageGenerationClient;
pub use crate::client::audio_generation::AudioGenerationClient;

pub use crate::client::{VerifyClient, VerifyError};

pub use crate::http_client::HttpLogConfigExt;