use std::io::Write;
use std::sync::Once;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::fmt;
static INIT: Once = Once::new();
pub const DEFAULT_LEVEL: &str = "warn";
pub const BODY_LOG_LIMIT: usize = 4096;
fn init_once(level: Option<&str>) {
INIT.call_once(|| {
if tracing::dispatcher::has_been_set() {
return;
}
let filter = build_filter(level);
let _ = fmt()
.with_env_filter(filter)
.with_writer(std::io::stderr)
.try_init();
});
}
pub fn init_logging(level: &str) {
init_once(Some(level));
}
pub fn auto_init_from_env() {
if INIT.is_completed() || tracing::dispatcher::has_been_set() {
return;
}
let env_present =
std::env::var_os("AIMUX_LOG").is_some() || std::env::var_os("AIMUX_LOG_LEVEL").is_some();
if env_present {
init_once(None);
}
}
#[must_use]
pub fn body_logging_enabled() -> bool {
std::env::var("AIMUX_LOG_BODY")
.map(|v| v == "1")
.unwrap_or(false)
}
fn build_filter(level: Option<&str>) -> EnvFilter {
if let Ok(directive) = std::env::var("AIMUX_LOG")
&& !directive.trim().is_empty()
{
return EnvFilter::try_new(directive).unwrap_or_else(|_| EnvFilter::new(DEFAULT_LEVEL));
}
let level = std::env::var("AIMUX_LOG_LEVEL")
.ok()
.filter(|s| !s.trim().is_empty())
.or_else(|| level.map(str::to_owned))
.unwrap_or_else(|| DEFAULT_LEVEL.to_owned());
EnvFilter::try_new(format!("aimux={level}")).unwrap_or_else(|_| EnvFilter::new(DEFAULT_LEVEL))
}
fn truncate(body: &str) -> &str {
if body.len() <= BODY_LOG_LIMIT {
return body;
}
let mut end = BODY_LOG_LIMIT;
while end > 0 && !body.is_char_boundary(end) {
end -= 1;
}
&body[..end]
}
#[must_use]
pub fn redact_body(body: &str) -> String {
let truncated = truncate(body);
match serde_json::from_str::<serde_json::Value>(truncated) {
Ok(value) => serde_json::to_string(&redact_error_context(value))
.unwrap_or_else(|_| truncated.to_owned()),
Err(_) => truncated.to_owned(),
}
}
const ERROR_CONTEXT_MAX_DEPTH: usize = 8;
const ERROR_CONTEXT_MAX_ITEMS: usize = 128;
const ERROR_CONTEXT_MAX_STRING: usize = 4096;
#[must_use]
pub fn redact_error_context(value: serde_json::Value) -> serde_json::Value {
redact_value(value, 0)
}
fn redact_value(value: serde_json::Value, depth: usize) -> serde_json::Value {
use serde_json::Value;
if depth >= ERROR_CONTEXT_MAX_DEPTH {
return Value::String("[MAX_DEPTH]".to_owned());
}
match value {
Value::Object(map) => {
let mut out = serde_json::Map::with_capacity(map.len().min(ERROR_CONTEXT_MAX_ITEMS));
for (index, (key, value)) in map.into_iter().enumerate() {
if index == ERROR_CONTEXT_MAX_ITEMS {
out.insert("__truncated__".into(), Value::Bool(true));
break;
}
if aimux_core::recording::is_sensitive_key(&key) {
out.insert(key, Value::String("[REDACTED]".to_owned()));
} else {
out.insert(key, redact_value(value, depth + 1));
}
}
Value::Object(out)
}
Value::Array(items) => {
let original_len = items.len();
let mut items: Vec<_> = items
.into_iter()
.take(ERROR_CONTEXT_MAX_ITEMS)
.map(|value| redact_value(value, depth + 1))
.collect();
if original_len > ERROR_CONTEXT_MAX_ITEMS {
items.push(serde_json::json!({
"__truncated_items__": original_len - ERROR_CONTEXT_MAX_ITEMS
}));
}
Value::Array(items)
}
Value::String(value) if value.len() > ERROR_CONTEXT_MAX_STRING => {
Value::String(format!("[STRING {} bytes]", value.len()))
}
other => other,
}
}
#[must_use]
pub fn redact_request_values(body: &crate::http::HttpBody) -> serde_json::Value {
match body {
crate::http::HttpBody::Json(value) => redact_error_context(value.clone()),
crate::http::HttpBody::Bytes(bytes, content_type) => serde_json::json!({
"content_type": content_type,
"length": bytes.len(),
}),
crate::http::HttpBody::Empty => serde_json::json!({}),
}
}
#[doc(hidden)]
pub struct CaptureWriter(pub std::sync::Arc<std::sync::Mutex<Vec<u8>>>);
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CaptureWriter {
type Writer = CaptureSink;
fn make_writer(&'a self) -> Self::Writer {
CaptureSink(self.0.clone())
}
}
#[doc(hidden)]
pub struct CaptureSink(pub std::sync::Arc<std::sync::Mutex<Vec<u8>>>);
impl Write for CaptureSink {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use serial_test::serial;
struct EnvGuard;
impl EnvGuard {
fn set(name: &str, value: &str) -> Self {
unsafe { std::env::set_var(name, value) };
Self
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
unsafe {
std::env::remove_var("AIMUX_LOG");
std::env::remove_var("AIMUX_LOG_LEVEL");
std::env::remove_var("AIMUX_LOG_BODY");
}
}
}
fn capture(
filter: EnvFilter,
) -> (
std::sync::Arc<std::sync::Mutex<Vec<u8>>>,
tracing::subscriber::DefaultGuard,
) {
let captured = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let subscriber = fmt()
.with_ansi(false)
.with_env_filter(filter)
.with_writer(CaptureWriter(captured.clone()))
.finish();
let guard = tracing::subscriber::set_default(subscriber);
(captured, guard)
}
fn captured_text(captured: &std::sync::Arc<std::sync::Mutex<Vec<u8>>>) -> String {
String::from_utf8(captured.lock().unwrap().clone()).unwrap()
}
#[test]
fn redact_json_body() {
let body = r#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"api_key":"sk-secret","nested":{"authorization":"Bearer xyz","token":"t"}}"#;
let out = redact_body(body);
assert!(!out.contains("sk-secret"), "api_key value leaked: {out}");
assert!(!out.contains("Bearer xyz"), "authorization leaked: {out}");
assert!(
out.contains("\"api_key\":\"[REDACTED]\""),
"api_key not redacted: {out}"
);
assert!(
out.contains("\"authorization\":\"[REDACTED]\""),
"authorization not redacted: {out}"
);
assert!(
out.contains("\"model\":\"gpt-4o\""),
"innocent field changed: {out}"
);
}
#[test]
fn redact_truncates_long_body() {
let body = "x".repeat(BODY_LOG_LIMIT + 500);
let out = redact_body(&body);
assert!(out.len() <= BODY_LOG_LIMIT);
}
#[test]
fn truncation_respects_utf8_char_boundaries() {
let body = format!("{}中中", "x".repeat(BODY_LOG_LIMIT - 1));
let out = redact_body(&body);
assert_eq!(out, "x".repeat(BODY_LOG_LIMIT - 1));
}
#[test]
fn redact_non_json_passthrough() {
let body = "not json at all";
assert_eq!(redact_body(body), body);
}
#[test]
#[serial]
fn filter_precedence_env_level() {
let _g = EnvGuard::set("AIMUX_LOG_LEVEL", "debug");
let (captured, _guard) = capture(build_filter(None));
tracing::debug!("visible-at-debug");
let out = captured_text(&captured);
assert!(
out.contains("visible-at-debug"),
"debug event not logged: {out}"
);
}
#[test]
#[serial]
fn filter_precedence_aimux_log_wins() {
let _g1 = EnvGuard::set("AIMUX_LOG", "aimux_provider_utils=error");
let _g2 = EnvGuard::set("AIMUX_LOG_LEVEL", "debug");
let (captured, _guard) = capture(build_filter(None));
tracing::debug!("debug-must-be-filtered");
tracing::error!("error-must-appear");
let out = captured_text(&captured);
assert!(
!out.contains("debug-must-be-filtered"),
"AIMUX_LOG=error should suppress debug: {out}"
);
assert!(
out.contains("error-must-appear"),
"error event missing: {out}"
);
}
#[test]
#[serial]
fn filter_defaults_to_warn() {
let (captured, _guard) = capture(build_filter(None));
tracing::debug!("debug-must-be-suppressed");
tracing::warn!("warn-must-appear");
let out = captured_text(&captured);
assert!(
!out.contains("debug-must-be-suppressed"),
"default level should suppress debug: {out}"
);
assert!(
out.contains("warn-must-appear"),
"warn event missing: {out}"
);
}
#[test]
#[serial]
fn init_is_idempotent() {
init_logging("debug");
init_logging("trace"); assert!(tracing::dispatcher::has_been_set());
}
}