diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 4bf41dc249c..64e16b512e1 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -105,6 +105,16 @@ jobs: with: python-version: "3.12" + - uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Install Python dependencies for the bridge tests + working-directory: . + run: | + uv sync --frozen --no-install-project + echo "PYTHONPATH=$PWD/.venv/lib/$(ls .venv/lib)/site-packages" >> "$GITHUB_ENV" + - run: rustup toolchain install --no-self-update - uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4427057a3b9..332d8b3dcf0 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2946,6 +2946,7 @@ name = "litellm-core-utils" version = "0.1.0" dependencies = [ "fancy-regex 0.19.2", + "litellm-tracing", "litellm-types", "rstest", "serde", @@ -3075,6 +3076,7 @@ dependencies = [ "aws-sdk-secretsmanager", "bytes", "criterion", + "fancy-regex 0.19.2", "futures-util", "litellm-auth", "litellm-auth-aws", @@ -3093,6 +3095,7 @@ dependencies = [ "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", + "litellm-host", "litellm-host-python", "litellm-http", "litellm-llms", @@ -3100,6 +3103,7 @@ dependencies = [ "litellm-secrets-aws", "litellm-secrets-types", "litellm-token-counter", + "litellm-tracing", "litellm-types", "pyo3", "pyo3-async-runtimes", @@ -3172,11 +3176,11 @@ dependencies = [ "litellm-auth-aws", "litellm-core-utils", "litellm-secrets-types", + "litellm-tracing", "rstest", "serde_json", "thiserror 2.0.19", "tokio", - "tracing", "veil", "wiremock", ] @@ -3208,6 +3212,7 @@ dependencies = [ "base64 0.22.1", "litellm-core-utils", "litellm-secrets-types", + "litellm-tracing", "moka", "percent-encoding", "reqwest 0.12.28", @@ -3216,7 +3221,6 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", - "tracing", "veil", "wiremock", ] @@ -3331,6 +3335,19 @@ dependencies = [ "tiktoken-rs", ] +[[package]] +name = "litellm-tracing" +version = "0.1.0" +dependencies = [ + "fancy-regex 0.19.2", + "percent-encoding", + "rstest", + "serde_json", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "litellm-types" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 65813a35214..9d05c8d2b98 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -9,6 +9,8 @@ license = "MIT" repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] +litellm-tracing = { path = "crates/tracing" } +tracing = "0.1" litellm-core = { path = "crates/core" } litellm-host = { path = "crates/host" } litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } diff --git a/litellm-rust/crates/core-utils/Cargo.toml b/litellm-rust/crates/core-utils/Cargo.toml index eb353bc060c..baf5dd16707 100644 --- a/litellm-rust/crates/core-utils/Cargo.toml +++ b/litellm-rust/crates/core-utils/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] fancy-regex.workspace = true +litellm-tracing.workspace = true litellm-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/core-utils/src/secret_redaction.rs b/litellm-rust/crates/core-utils/src/secret_redaction.rs index e3caee4799a..14a88eae288 100644 --- a/litellm-rust/crates/core-utils/src/secret_redaction.rs +++ b/litellm-rust/crates/core-utils/src/secret_redaction.rs @@ -1,109 +1 @@ -use fancy_regex::Regex; - -pub const REDACTED: &str = "REDACTED"; - -const DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH: usize = 16; - -fn minimum_custom_key_length() -> usize { - std::env::var("MINIMUM_CUSTOM_KEY_LENGTH") - .ok() - .and_then(|value| value.trim().parse().ok()) - .unwrap_or(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH) -} - -fn secret_patterns(minimum_custom_key_length: usize) -> String { - let sk_suffix_length = minimum_custom_key_length.saturating_sub("sk-".len()); - [ - r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----", - r"\bya29\.[A-Za-z0-9_.~+/-]+", - r#"(?:client_secret|azure_password|azure_username)\s+[^\s,'"})\]{}>]+"#, - r"(?:AKIA|ASIA)[0-9A-Z]{16}", - r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*", - r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", - &format!(r"sk-[A-Za-z0-9\-_]{{{sk_suffix_length},}}"), - r#"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))=[^\s&'"]+"#, - r#"(?:api[_-]?key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]{8,}"#, - r#"(?:x-api-key|api-key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, - r"x-ak-[A-Za-z0-9\-_]{20,}", - r"AIza[0-9A-Za-z\-_]{35}", - r#"(?<=[?&])key=[^\s&'"]{8,}"#, - r#"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, - r#"(?<=://)[^\s'":]{0,4096}:[^\s'"]{1,4096}(?=@)"#, - r"dapi[0-9a-f]{32}", - r#"litellm\.[A-Za-z0-9_]*_key['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, - r#"private_key['"]?\s*[:=]\s*['"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'"})\]{}>]+)"#, - concat!( - r"(?:master_key|xai_key|database_url|db_url|connection_string|", - r"aws_secret_access_key|aws_session_token|aws_access_key_id|", - r"signing_key|encryption_key|", - r"auth_token|access_token|refresh_token|", - r"slack_webhook_url|webhook_url|", - r"database_connection_string|", - r"huggingface_token|jwt_secret)", - r#"['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, - ), - r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", - r"(?<=[?&])sig=[A-Za-z0-9%+/=]+", - r#"\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"#, - ] - .join("|") -} - -/// Python's `_ENABLE_SECRET_REDACTION` pattern set, compiled once per configuration. -#[derive(Clone, Debug)] -pub struct SecretRedactor { - pattern: Regex, -} - -impl SecretRedactor { - pub fn new(minimum_custom_key_length: usize) -> Self { - let pattern = Regex::new(&format!( - "(?i){}", - secret_patterns(minimum_custom_key_length) - )) - .expect("secret redaction patterns compile"); - Self { pattern } - } - - /// `None` when `LITELLM_DISABLE_REDACT_SECRETS` turns redaction off. - pub fn from_env() -> Option { - let disabled = std::env::var("LITELLM_DISABLE_REDACT_SECRETS") - .is_ok_and(|value| value.eq_ignore_ascii_case("true")); - (!disabled).then(|| Self::new(minimum_custom_key_length())) - } - - pub fn redact(&self, value: &str) -> String { - self.pattern.replace_all(value, REDACTED).into_owned() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[rstest::rstest] - #[case::bearer("auth failed: Bearer abcdefghijklmnop", "auth failed: REDACTED")] - #[case::sk_key("key sk-abcdefghijklmnopqrstuvwxyz rejected", "key REDACTED rejected")] - #[case::short_sk_key_is_kept("sk-abc", "sk-abc")] - #[case::query_param("GET /v1?api_key=secret123&x=1", "GET /v1?REDACTED&x=1")] - #[case::dict_repr("{'api_key': 'abcdefghij'}", "{'REDACTED'}")] - #[case::url_credentials("postgres://user:pass@host/db", "postgres://REDACTED@host/db")] - #[case::case_insensitive("BEARER ABCDEFGHIJKLMNOP", "REDACTED")] - #[case::aws_key("AKIAABCDEFGHIJKLMNOP", "REDACTED")] - #[case::sas_signature("https://x.blob/a?sv=1&sig=abc%2B=", "https://x.blob/a?sv=1&REDACTED")] - #[case::password_needs_word_boundary("db_password=hunter2", "REDACTED")] - #[case::plain_text_is_kept(r#"{"message": "rejected"}"#, r#"{"message": "rejected"}"#)] - fn redacts_the_same_spans_as_the_python_patterns(#[case] input: &str, #[case] expected: &str) { - assert_eq!( - SecretRedactor::new(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH).redact(input), - expected - ); - } - - #[test] - fn sk_threshold_follows_the_minimum_custom_key_length() { - let redactor = SecretRedactor::new(8); - assert_eq!(redactor.redact("sk-abcde"), REDACTED); - assert_eq!(redactor.redact("sk-abcd"), "sk-abcd"); - } -} +pub use litellm_tracing::{REDACTED, SecretRedactor}; diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index dd0975219e4..85e6b9588dd 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,9 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +fancy-regex.workspace = true +litellm-tracing.workspace = true +litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/activation.rs b/litellm-rust/crates/python-bridge/src/cache/activation.rs index 8d2c340dec8..f77032c579d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/activation.rs +++ b/litellm-rust/crates/python-bridge/src/cache/activation.rs @@ -1,6 +1,7 @@ +use crate::logger::run_sync_value; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_redis_semantic::RedisSemanticConfig; -use litellm_host_python::{release_gil, run_sync_value}; +use litellm_host_python::release_gil; use litellm_http::ClientVariant; use pyo3::prelude::*; diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 56f1bff2853..fcf71e10427 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -1,5 +1,6 @@ +use crate::logger::run_async; use litellm_cache_response::PartialHits; -use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use litellm_host_python::{ExecutionStep, from_py, release_gil, to_py}; use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyValueError}, diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index dff8a771a23..e14916b25c6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,10 +1,11 @@ +use crate::logger::run_sync_value; use litellm_auth_aws::AwsAuthConfig; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization}; use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; -use litellm_host_python::{release_gil, run_sync_value}; +use litellm_host_python::release_gil; use litellm_http::ClientVariant; use pyo3::{ PyTraverseError, PyVisit, diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 3c260e70843..0e279046812 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -470,7 +470,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - litellm_host_python::run_async( + crate::logger::run_async( py, async move { service @@ -495,7 +495,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - litellm_host_python::run_async( + crate::logger::run_async( py, async move { service.async_lookup(&request, now()).await }, super::cache_error, @@ -550,7 +550,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - litellm_host_python::run_async( + crate::logger::run_async( py, async move { service.async_store(&request, response, now()).await }, super::cache_error, @@ -619,7 +619,7 @@ impl NativeResponseCache { match self { Self::Exact(_) | Self::QdrantSemantic(_) => { let service = self.clone(); - litellm_host_python::run_async( + crate::logger::run_async( py, async move { service.async_store_batch(entries, now()).await }, super::cache_error, diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs index bea58b88885..eb67f5594e2 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -1,7 +1,8 @@ +use crate::logger::run_async; use std::{collections::VecDeque, time::Duration}; use litellm_cache::Error; -use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep}; use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyException, PyRuntimeError}, diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 2515b409c54..3dad3447f45 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -102,7 +102,7 @@ pub(crate) fn call_config( .without_missing_files(&|path: &Path| path.exists()); let resolution = Resolution::from(&settings); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { - PythonSettings::warn(py, &unsupported.to_string())?; + crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}")); } Ok(resolution.config) } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index b1fc5244d6f..76de2ae2b7b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -4,6 +4,7 @@ mod credentials; mod diagnostics; mod errors; mod http; +mod logger; mod marshal; mod python_settings; mod routes; @@ -25,6 +26,8 @@ mod _native { #[pymodule_export] use crate::errors::{RustBridgeDeclined, RustUpstreamError}; #[pymodule_export] + use crate::logger::NativeDiagnosticProcessor; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{ @@ -87,6 +90,7 @@ mod tests { "chat_completions", "achat_completions", "ResponsesWebSocketConnection", + "NativeDiagnosticProcessor", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/logger/execution.rs b/litellm-rust/crates/python-bridge/src/logger/execution.rs new file mode 100644 index 00000000000..c8d5c0023e3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/logger/execution.rs @@ -0,0 +1,46 @@ +use std::future::Future; + +use pyo3::prelude::*; +use serde::Serialize; + +pub(crate) fn run_sync( + py: Python<'_>, + future: F, + map_error: fn(E) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, +{ + litellm_host_python::run_sync(py, super::capture(py).instrument(future), map_error) +} + +pub(crate) fn run_async( + py: Python<'_>, + future: F, + map_error: fn(E) -> PyErr, +) -> PyResult> +where + T: Serialize + Send + 'static, + E: Send + 'static, + F: Future> + Send + 'static, +{ + litellm_host_python::run_async(py, super::capture(py).instrument(future), map_error) +} + +pub(crate) fn run_sync_value(py: Python<'_>, future: F) -> PyResult +where + T: Send + 'static, + F: Future> + Send + 'static, +{ + litellm_host_python::run_sync_value(py, super::capture(py).instrument(future)) +} + +pub(crate) fn run_async_value(py: Python<'_>, future: F) -> PyResult> +where + T: for<'py> IntoPyObject<'py> + Send + 'static, + F: Future> + Send + 'static, +{ + litellm_host_python::run_async_value(py, super::capture(py).instrument(future)) +} diff --git a/litellm-rust/crates/python-bridge/src/logger/machine.rs b/litellm-rust/crates/python-bridge/src/logger/machine.rs new file mode 100644 index 00000000000..7234308e67e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/logger/machine.rs @@ -0,0 +1,41 @@ +use std::sync::OnceLock; + +use litellm_host::{ + host::HostResult, + machine::{HostFailure, Interrupted, Machine, Step}, + route::Route, +}; +use litellm_tracing::Logger; +use pyo3::Python; + +pub(crate) struct LoggedMachine { + machine: M, + logger: OnceLock, +} + +impl LoggedMachine { + pub(crate) fn new(machine: M) -> Self { + Self { + machine, + logger: OnceLock::new(), + } + } +} + +impl Machine for LoggedMachine { + type Route = M::Route; + type Complete = M::Complete; + + fn resume(&mut self, result: Option>) -> Step<'_, Self> { + let logger = self.logger.get_or_init(|| Python::attach(super::capture)); + Box::pin(logger.instrument(logger.scope(|| self.machine.resume(result)))) + } + + fn interrupt( + &mut self, + failure: HostFailure<::Error>, + ) -> Interrupted<'_, Self> { + let logger = self.logger.get_or_init(|| Python::attach(super::capture)); + Box::pin(logger.instrument(logger.scope(|| self.machine.interrupt(failure)))) + } +} diff --git a/litellm-rust/crates/python-bridge/src/logger/mod.rs b/litellm-rust/crates/python-bridge/src/logger/mod.rs new file mode 100644 index 00000000000..6421fe1d554 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/logger/mod.rs @@ -0,0 +1,166 @@ +mod execution; +mod machine; + +pub(crate) use execution::{run_async, run_async_value, run_sync, run_sync_value}; +pub(crate) use machine::LoggedMachine; + +use litellm_host_python::Pythonized; +use litellm_tracing::{DiagnosticInput, Level, Logger, Metadata, Policy, Processor, Record, Sink}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; + +const MODULE: &str = "litellm.rust_bridge.logger"; +type NativeDiagnosticOutput = (String, Option, Option, Vec, bool); + +#[pyclass] +pub(crate) struct NativeDiagnosticProcessor { + inner: Processor, +} + +#[pymethods] +impl NativeDiagnosticProcessor { + #[new] + fn new(minimum_custom_key_length: usize) -> Self { + Self { + inner: Processor::new(minimum_custom_key_length), + } + } + + fn redact_text(&self, text: &str) -> PyResult { + self.inner.redact_text(text).map_err(processing_error) + } + + fn redact_structured_text(&self, key: Option<&str>, text: &str) -> PyResult { + self.inner + .redact_structured_text(key, text) + .map_err(processing_error) + } + + fn redact_client_message(&self, text: &str) -> PyResult { + self.inner + .redact_client_message(text) + .map_err(processing_error) + } + + #[pyo3(signature = (message, exception, stack, leaves, policy))] + fn process_diagnostic( + &self, + message: String, + exception: Option, + stack: Option, + leaves: Vec<(Option, String)>, + policy: (bool, i64, i64), + ) -> PyResult { + let input = DiagnosticInput { + message, + exception, + stack, + leaves, + }; + let policy = Policy { + redact: policy.0, + base64_limit: policy.1, + text_limit: policy.2, + }; + self.inner + .process_diagnostic(&input, policy) + .map(|output| { + ( + output.message, + output.exception, + output.stack, + output.leaves, + output.changed, + ) + }) + .map_err(processing_error) + } + + fn scrub_access_arguments(&self, arguments: Vec) -> PyResult> { + self.inner + .scrub_access_arguments(&arguments) + .map_err(processing_error) + } +} + +fn processing_error(_: fancy_regex::Error) -> PyErr { + PyRuntimeError::new_err("diagnostic processing failed") +} + +struct PythonSink { + correlation: (String, String), +} + +fn level(level: &Level) -> u8 { + match *level { + Level::ERROR => 40, + Level::WARN => 30, + Level::INFO => 20, + Level::DEBUG | Level::TRACE => 10, + } +} + +fn report(py: Python<'_>, result: PyResult) -> T { + match result { + Ok(value) => value, + Err(error) => { + error.write_unraisable(py, None); + T::default() + } + } +} + +impl Sink for PythonSink { + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + if !metadata.target().starts_with("litellm_") && !metadata.target().starts_with("_native::") + { + return false; + } + Python::try_attach(|py| { + report( + py, + py.import(MODULE) + .and_then(|module| module.call_method1("enabled", (level(metadata.level()),))) + .and_then(|enabled| enabled.extract()), + ) + }) + .unwrap_or(false) + } + + fn emit(&self, record: &Record) { + Python::try_attach(|py| { + report( + py, + py.import(MODULE).and_then(|module| { + module + .call_method1( + "emit", + ( + level(record.metadata.level()), + &record.message, + record.metadata.file().unwrap_or_default(), + record.metadata.line().unwrap_or_default(), + record.metadata.target(), + Pythonized(&record.fields), + (&self.correlation.0, &self.correlation.1), + ), + ) + .map(|_| ()) + }), + ); + }); + } +} + +pub(crate) fn capture(py: Python<'_>) -> Logger { + report( + py, + py.import(MODULE) + .and_then(|module| module.call_method0("context")) + .and_then(|value| value.extract()) + .map(|correlation| Logger::new(PythonSink { correlation })), + ) +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs new file mode 100644 index 00000000000..9312d4c187c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -0,0 +1,295 @@ +use std::{process::Command, task::Poll}; + +use litellm_host::{ + host::HostResult, + machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, + route::Route, +}; + +use pyo3::{prelude::*, types::PyDict}; + +struct DiagnosticMachine; + +impl Route for DiagnosticMachine { + type Response = (); + type Error = String; + type Op = (); + type OpResult = (); + type Chunk = (); + type StreamHead = (); +} + +impl Machine for DiagnosticMachine { + type Route = Self; + type Complete = (); + + fn resume(&mut self, _: Option>) -> Step<'_, Self> { + litellm_tracing::warn!("machine started"); + Box::pin(async { + tokio::task::yield_now().await; + litellm_tracing::warn!("machine warning"); + Ok(MachineStep::Complete(())) + }) + } + + fn interrupt(&mut self, _: HostFailure) -> Interrupted<'_, Self> { + Box::pin(async { + litellm_tracing::warn!("machine interrupted"); + Ok(()) + }) + } +} + +#[pyfunction] +fn machine_warning(py: Python<'_>) -> PyResult> { + let mut machine = super::LoggedMachine::new(DiagnosticMachine); + let mut future = Box::pin(async move { + machine + .resume(None) + .await + .map_err(pyo3::exceptions::PyValueError::new_err)?; + machine + .interrupt(HostFailure::Error("stop".into())) + .await + .map_err(pyo3::exceptions::PyValueError::new_err) + }); + assert!(matches!( + litellm_host_python::poll_async_value(py, future.as_mut())?, + Poll::Pending + )); + litellm_host_python::run_async_value(py, future) +} + +#[pyfunction] +fn warning(py: Python<'_>) { + super::capture(py).scope(|| { + litellm_tracing::warn!(attempt = 3, retry = true, "native warning"); + }); +} + +#[pyfunction] +fn levels(py: Python<'_>) { + super::capture(py).scope(|| { + litellm_tracing::trace!("trace"); + litellm_tracing::debug!("debug"); + litellm_tracing::info!("info"); + litellm_tracing::warn!("warn"); + litellm_tracing::error!("error"); + litellm_tracing::warn!(target: "unrelated_transport", "private wire data"); + }); +} + +#[pyfunction] +fn asynchronous_warning(py: Python<'_>) -> PyResult> { + super::run_async_value(py, async { + tokio::task::yield_now().await; + litellm_tracing::warn!("async warning"); + Ok(()) + }) +} + +#[pyfunction] +fn synchronous_warning(py: Python<'_>) -> PyResult<()> { + super::run_sync_value(py, async { + tokio::task::yield_now().await; + litellm_tracing::warn!("sync warning"); + Ok(()) + }) +} + +#[pyfunction] +fn synchronous_failure(py: Python<'_>) -> PyResult<()> { + super::run_sync_value(py, async { + litellm_tracing::warn!("failure diagnostic"); + Err(pyo3::exceptions::PyValueError::new_err("request failed")) + }) +} + +#[pyfunction] +fn http_warning(py: Python<'_>) -> PyResult<()> { + crate::http::call_config(py, &PyDict::new(py), false).map(|_| ()) +} + +#[test] +fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication() { + if std::env::var_os("LITELLM_LOGGER_TEST_PROCESS").is_none() { + let output = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + std::thread::current().name().unwrap(), + "--nocapture", + ]) + .env("LITELLM_LOGGER_TEST_PROCESS", "1") + .output() + .unwrap(); + assert!( + output.status.success(), + "{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + return; + } + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals + .set_item( + "repo_root", + concat!(env!("CARGO_MANIFEST_DIR"), "/../../.."), + ) + .unwrap(); + locals + .set_item( + "machine_warning", + wrap_pyfunction!(machine_warning, py).unwrap(), + ) + .unwrap(); + locals + .set_item( + "synchronous_failure", + wrap_pyfunction!(synchronous_failure, py).unwrap(), + ) + .unwrap(); + locals + .set_item("levels", wrap_pyfunction!(levels, py).unwrap()) + .unwrap(); + locals + .set_item("warning", wrap_pyfunction!(warning, py).unwrap()) + .unwrap(); + locals + .set_item( + "asynchronous_warning", + wrap_pyfunction!(asynchronous_warning, py).unwrap(), + ) + .unwrap(); + locals + .set_item( + "synchronous_warning", + wrap_pyfunction!(synchronous_warning, py).unwrap(), + ) + .unwrap(); + locals + .set_item("http_warning", wrap_pyfunction!(http_warning, py).unwrap()) + .unwrap(); + let importable = py + .eval( + c"__import__('importlib.util', fromlist=['util']).find_spec('dotenv') is not None", + Some(&locals), + Some(&locals), + ) + .unwrap() + .is_truthy() + .unwrap(); + if !importable { + eprintln!("SKIP: litellm package dependencies are not importable in this interpreter"); + return; + } + py.run(c" +import asyncio +import logging +import sys +sys.path.insert(0, repo_root) +import litellm +from litellm._logging import verbose_logger, session_id_var, trace_id_var + +class Capture(logging.Handler): + def __init__(self): + super().__init__() + self.records = [] + def emit(self, record): + self.records.append(record) + warning() + +class Broken(logging.Handler): + def emit(self, record): + raise ValueError('handler failed') + +capture = Capture() +old_handlers = verbose_logger.handlers +old_level = verbose_logger.level +old_correlation = litellm.request_correlation_in_logs +old_curve = litellm.ssl_ecdh_curve +old_unraisable = sys.unraisablehook +failures = [] +try: + verbose_logger.handlers = [capture] + litellm.request_correlation_in_logs = True + verbose_logger.setLevel(logging.ERROR) + warning() + assert capture.records == [] + verbose_logger.setLevel(logging.WARNING) + warning() + assert len(capture.records) == 1 + record = capture.records[0] + assert record.getMessage() == 'native warning' + assert record.levelno == logging.WARNING + assert record.rust_fields == {'attempt': 3, 'retry': True} + assert record.pathname.endswith('logger/tests.rs') + assert record.lineno > 0 + assert record.rust_target.endswith('logger::tests') + verbose_logger.setLevel(logging.ERROR) + warning() + assert len(capture.records) == 1 + verbose_logger.setLevel(logging.WARNING) + + async def request(name): + session = session_id_var.set(name) + trace = trace_id_var.set('trace-' + name) + try: + await asynchronous_warning() + await machine_warning() + synchronous_warning() + assert session_id_var.get() == name + assert trace_id_var.get() == 'trace-' + name + finally: + trace_id_var.reset(trace) + session_id_var.reset(session) + + async def concurrent(): + await asyncio.gather(request('first'), request('second')) + + asyncio.run(concurrent()) + assert sorted((r.getMessage(), r.session_id, r.trace_id) for r in capture.records[1:]) == sorted( + (message, name, 'trace-' + name) + for name in ('first', 'second') + for message in ('async warning', 'sync warning', 'machine started', 'machine warning', 'machine interrupted') + ) + + verbose_logger.setLevel(logging.DEBUG) + before_levels = len(capture.records) + levels() + assert [(r.getMessage(), r.levelno) for r in capture.records[before_levels:]] == [ + ('trace', logging.DEBUG), ('debug', logging.DEBUG), ('info', logging.INFO), + ('warn', logging.WARNING), ('error', logging.ERROR), + ] + + before = len(capture.records) + litellm.ssl_ecdh_curve = 'logger-test-unsupported-curve' + http_warning() + http_warning() + assert len(capture.records) == before + 1 + assert 'logger-test-unsupported-curve' in capture.records[-1].getMessage() + assert capture.records[-1].pathname.endswith('http.rs') + + verbose_logger.handlers = [Broken()] + sys.unraisablehook = failures.append + warning() + assert len(failures) == 1 + assert str(failures[0].exc_value) == 'handler failed' + try: + synchronous_failure() + except ValueError as error: + assert str(error) == 'request failed' + else: + raise AssertionError('request failure was lost') + assert len(failures) == 2 +finally: + sys.unraisablehook = old_unraisable + verbose_logger.handlers = old_handlers + verbose_logger.setLevel(old_level) + litellm.request_correlation_in_logs = old_correlation + litellm.ssl_ecdh_curve = old_curve +", Some(&locals), Some(&locals)).unwrap(); + }); +} diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 111ac3bc259..abf664b795d 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -44,11 +44,6 @@ impl PythonSettings { pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { Snapshot { group: self, value } } - - pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> { - py.import(MODULE)?.getattr("warn")?.call1((message,))?; - Ok(()) - } } #[cfg(test)] diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index d63e9a1feaf..dec4dcea21c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -1,7 +1,8 @@ +use crate::logger::{run_async, run_sync}; use litellm_core::audio_transcription::{ Error, audio_transcription as run_audio_transcription, types::AudioTranscriptionRequest, }; -use litellm_host_python::{from_py_argument, run_async, run_sync}; +use litellm_host_python::from_py_argument; use pyo3::prelude::*; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 049a507dcdc..1fa2ca00c42 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -1,8 +1,9 @@ +use crate::logger::{run_async, run_sync}; use litellm_core::chat_completions::{ Error, chat_completions as run_chat_completions, chat_completions_decline_reason, types::ChatCompletionsRequest, }; -use litellm_host_python::{from_py_argument, run_async, run_sync}; +use litellm_host_python::from_py_argument; use litellm_types::utils::ChatCompletionsResponse; use pyo3::prelude::*; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index b606293f79f..804d883e9ae 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -43,7 +43,7 @@ fn run_messages( py, SURFACE, PublicCall::capture(&request, &args, &kwargs)?, - messages_machine(), + crate::logger::LoggedMachine::new(messages_machine()), MessagesRouteHost::new(request.unbind()), asynchronous, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 2dca6da66cd..ca5abf1f80e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -73,7 +73,7 @@ fn run_ocr( py, if asynchronous { ASYNC_SURFACE } else { SURFACE }, PublicCall::capture(&request, &args, &kwargs)?, - ocr_machine(client), + crate::logger::LoggedMachine::new(ocr_machine(client)), OcrRouteHost::new(request.unbind()), asynchronous, ) diff --git a/litellm-rust/crates/python-bridge/src/routes/responses.rs b/litellm-rust/crates/python-bridge/src/routes/responses.rs index 2e7e8fcbc21..ffbb945c415 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses.rs @@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection { ) -> PyResult> { let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); - litellm_host_python::run_async_value(py, async move { + crate::logger::run_async_value(py, async move { let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) .await .map_err(responses_error_to_pyerr)?; @@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection { fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult> { let inner = self.inner.clone(); - litellm_host_python::run_async_value(py, async move { + crate::logger::run_async_value(py, async move { inner .send_text(text) .await @@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection { fn recv_text<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); - litellm_host_python::run_async_value(py, async move { + crate::logger::run_async_value(py, async move { inner.recv_text().await.map_err(responses_error_to_pyerr) }) } fn close<'py>(&self, py: Python<'py>) -> PyResult> { let inner = self.inner.clone(); - litellm_host_python::run_async_value(py, async move { + crate::logger::run_async_value(py, async move { inner.close().await.map_err(responses_error_to_pyerr) }) } diff --git a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs index 168c4883b0b..21589aa3fe9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/token_counter.rs +++ b/litellm-rust/crates/python-bridge/src/routes/token_counter.rs @@ -1,7 +1,8 @@ +use crate::logger::run_async; use std::sync::Arc; use std::{num::NonZero, thread::available_parallelism}; -use litellm_host_python::{enter_native, run_async}; +use litellm_host_python::enter_native; use litellm_token_counter::{ CountableRequest, Error, InputTokenCount, TokenCounter as CoreTokenCounter, }; diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml index b1dc5b33cda..9c926bb441b 100644 --- a/litellm-rust/crates/secrets-aws/Cargo.toml +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -11,7 +11,7 @@ litellm-secrets-types.workspace = true litellm-core-utils.workspace = true serde_json.workspace = true thiserror.workspace = true -tracing = "0.1" +litellm-tracing.workspace = true veil.workspace = true aws-sdk-kms = "1.120.0" aws-sdk-secretsmanager = "1.117.0" diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 76c61534861..5658b6d7e11 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -215,7 +215,7 @@ impl AwsSecretsManagerV2 { .await .is_err() { - tracing::warn!("secret created but replication failed"); + litellm_tracing::warn!("secret created but replication failed"); } Ok(response) } diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index 3c1159c40be..c1613f96403 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -14,7 +14,7 @@ reqwest.workspace = true serde_json.workspace = true thiserror.workspace = true veil.workspace = true -tracing = "0.1" +litellm-tracing.workspace = true percent-encoding = "2.3" tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index 80e6fb12f9e..dd36cf183d8 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -92,7 +92,7 @@ impl CyberArkSecretManager { .unwrap_or(true); let mut builder = reqwest::Client::builder(); if !verify { - tracing::warn!( + litellm_tracing::warn!( "CyberArk SSL verification is disabled. This is insecure and should only be used for testing with self-signed certificates." ); builder = builder.danger_accept_invalid_certs(true); @@ -259,11 +259,13 @@ impl CyberArkSecretManager { .endpoint .join(&format!("policies/{}/policy/root", self.account)); let Ok(policy_url) = policy_url else { - tracing::warn!("Could not build CyberArk policy endpoint"); + litellm_tracing::warn!("Could not build CyberArk policy endpoint"); return; }; let Ok(authorization) = self.authorization_header(context).await else { - tracing::warn!("Could not authenticate while ensuring CyberArk variable exists"); + litellm_tracing::warn!( + "Could not authenticate while ensuring CyberArk variable exists" + ); return; }; let body = format!( @@ -288,19 +290,19 @@ impl CyberArkSecretManager { reqwest::StatusCode::CONFLICT | reqwest::StatusCode::UNPROCESSABLE_ENTITY ) => { - tracing::debug!( + litellm_tracing::debug!( "CyberArk variable policy already exists or conflicts: {}", response.status() ); } Ok(response) => { - tracing::warn!( + litellm_tracing::warn!( "Could not ensure CyberArk variable exists: {}", response.status() ); } Err(error) => { - tracing::warn!("Error ensuring CyberArk variable exists: {error}"); + litellm_tracing::warn!("Error ensuring CyberArk variable exists: {error}"); } } } @@ -324,7 +326,7 @@ impl CyberArkSecretManager { _recovery_window_in_days: Option, _context: &SecretOperationContext, ) -> Result { - tracing::warn!( + litellm_tracing::warn!( "CyberArk Conjur does not support direct secret deletion. Secrets must be removed through policy updates." ); self.secrets.invalidate(name).await; diff --git a/litellm-rust/crates/tracing/Cargo.toml b/litellm-rust/crates/tracing/Cargo.toml new file mode 100644 index 00000000000..41ad20afb3e --- /dev/null +++ b/litellm-rust/crates/tracing/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-tracing" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +fancy-regex.workspace = true +percent-encoding.workspace = true +serde_json.workspace = true +tracing.workspace = true +tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/tracing/README.md b/litellm-rust/crates/tracing/README.md new file mode 100644 index 00000000000..7f086067f3e --- /dev/null +++ b/litellm-rust/crates/tracing/README.md @@ -0,0 +1,26 @@ +# Native diagnostic tracing + +`litellm-tracing` connects standard `tracing` events to a host-provided `Sink`. It has no Python dependency and does not install a global subscriber + +Use the exported `debug!`, `info!`, `warn!`, and `error!` macros in native code. A host creates a `Logger` with its sink, uses `scope` for synchronous operations, and wraps futures with `instrument`. Instrument spawned futures explicitly because thread-local subscribers do not automatically follow spawned work + +Bindings implement `litellm_tracing::Sink` to connect events to their host runtime: + +```rust +pub trait Sink: Send + Sync + 'static { + fn enabled(&self, metadata: &Metadata<'_>) -> bool; + fn emit(&self, record: &Record); +} +``` + +Pass the implementation to `litellm_tracing::Logger::new(sink)`, then call `logger.scope(|| litellm_tracing::info!(attempt = 1, "request started"))`. The sink owns host access, level mapping, correlation capture, and delivery failures. `enabled` runs before event fields are evaluated or formatted. `emit` borrows a record; an adapter that queues delivery must copy the data it needs into an owned value + +Records retain event metadata, the message, and typed event fields. Sink filtering runs for each event so runtime level changes take effect. Logging from inside a sink is suppressed to prevent recursion + +The Python bridge scopes native execution to a sink that uses LiteLLM's existing Python logger. It preserves request correlation, redacts before delivering to handlers, maps Rust trace events to Python debug, and reports handler failures through `sys.unraisablehook`. It accepts LiteLLM targets only, keeping dependency wire diagnostics out of the application logger + +Python consumers continue using `litellm._logging` and its existing loggers, filters, formatters, and context setters. Catalog dispatch selects the processing backend for both Python and native diagnostics. The pure `Processor` takes explicit settings and never emits events + +A future Node bridge can implement the same sink with runtime-specific delivery and expose the same processor through N-API. Node callback scheduling, queue limits, and shutdown belong in that bridge; this crate has no interpreter handles or output queue + +This is diagnostic logging. Request lifecycle hooks and `CustomLogger` dispatch remain separate diff --git a/litellm-rust/crates/tracing/src/lib.rs b/litellm-rust/crates/tracing/src/lib.rs new file mode 100644 index 00000000000..47f97d6db27 --- /dev/null +++ b/litellm-rust/crates/tracing/src/lib.rs @@ -0,0 +1,146 @@ +use std::{ + cell::Cell, + fmt, + future::{Future, poll_fn}, + pin::pin, +}; + +use serde_json::{Map, Value}; +use tracing::{ + Dispatch, Event, Subscriber, + field::{Field, Visit}, + subscriber::Interest, +}; +use tracing_subscriber::{Layer, Registry, layer::Context, prelude::*}; + +mod processing; +mod redaction; + +pub use processing::{DiagnosticInput, DiagnosticOutput, Policy, Processor}; +pub use redaction::{REDACTED, SecretRedactor}; +pub use tracing::{Level, Metadata, debug, error, info, trace, warn}; + +pub trait Sink: Send + Sync + 'static { + fn enabled(&self, metadata: &Metadata<'_>) -> bool; + fn emit(&self, record: &Record); +} + +#[derive(Debug)] +pub struct Record { + pub metadata: &'static Metadata<'static>, + pub message: String, + pub fields: Map, +} + +#[derive(Clone, Default)] +pub struct Logger { + dispatch: Dispatch, +} + +impl Logger { + pub fn new(sink: impl Sink) -> Self { + Self { + dispatch: Dispatch::new(Registry::default().with(Output(sink))), + } + } + + pub fn scope(&self, operation: impl FnOnce() -> T) -> T { + if EMITTING.get() { + return operation(); + } + tracing::dispatcher::with_default(&self.dispatch, operation) + } + + pub fn instrument(&self, future: F) -> impl Future + use { + let logger = self.clone(); + async move { + let mut future = pin!(future); + poll_fn(|context| logger.scope(|| future.as_mut().poll(context))).await + } + } +} + +thread_local! { + static EMITTING: Cell = const { Cell::new(false) }; +} + +struct Emitting; + +impl Emitting { + fn enter() -> Option { + EMITTING.with(|active| (!active.replace(true)).then_some(Self)) + } +} + +impl Drop for Emitting { + fn drop(&mut self) { + EMITTING.set(false); + } +} + +struct Output(S); + +impl Layer for Output { + fn register_callsite(&self, _: &'static Metadata<'static>) -> Interest { + Interest::sometimes() + } + + fn enabled(&self, metadata: &Metadata<'_>, _: Context<'_, R>) -> bool { + let Some(_guard) = Emitting::enter() else { + return false; + }; + self.0.enabled(metadata) + } + + fn on_event(&self, event: &Event<'_>, _: Context<'_, R>) { + let Some(_guard) = Emitting::enter() else { + return; + }; + let mut record = Record { + metadata: event.metadata(), + message: String::new(), + fields: Map::new(), + }; + event.record(&mut record); + self.0.emit(&record); + } +} + +impl Record { + fn field(&mut self, field: &Field, value: Value) { + if field.name() == "message" { + self.message = match value { + Value::String(message) => message, + value => value.to_string(), + }; + } else { + self.fields.insert(field.name().to_owned(), value); + } + } +} + +impl Visit for Record { + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { + self.field(field, format!("{value:?}").into()); + } + + fn record_str(&mut self, field: &Field, value: &str) { + self.field(field, value.into()); + } + + fn record_bool(&mut self, field: &Field, value: bool) { + self.field(field, value.into()); + } + + fn record_i64(&mut self, field: &Field, value: i64) { + self.field(field, value.into()); + } + + fn record_u64(&mut self, field: &Field, value: u64) { + self.field(field, value.into()); + } + + fn record_f64(&mut self, field: &Field, value: f64) { + self.field(field, value.into()); + } +} diff --git a/litellm-rust/crates/tracing/src/processing.rs b/litellm-rust/crates/tracing/src/processing.rs new file mode 100644 index 00000000000..29f5f6c71ea --- /dev/null +++ b/litellm-rust/crates/tracing/src/processing.rs @@ -0,0 +1,333 @@ +use fancy_regex::Result; +use percent_encoding::percent_decode_str; + +use crate::{REDACTED, SecretRedactor}; + +#[derive(Clone, Copy, Debug)] +pub struct Policy { + pub redact: bool, + pub base64_limit: i64, + pub text_limit: i64, +} + +#[derive(Clone, Debug)] +pub struct DiagnosticInput { + pub message: String, + pub exception: Option, + pub stack: Option, + pub leaves: Vec<(Option, String)>, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DiagnosticOutput { + pub message: String, + pub exception: Option, + pub stack: Option, + pub leaves: Vec, + pub changed: bool, +} + +pub struct Processor { + redactor: SecretRedactor, +} + +impl Processor { + pub fn new(minimum_custom_key_length: usize) -> Self { + Self { + redactor: SecretRedactor::new(minimum_custom_key_length), + } + } + + pub fn redact_text(&self, text: &str) -> Result { + self.redactor.try_redact(text) + } + + pub fn redact_structured_text(&self, key: Option<&str>, text: &str) -> Result { + self.redactor.try_redact_structured(key, text) + } + + pub fn redact_client_message(&self, text: &str) -> Result { + self.redactor.try_redact_internal(text) + } + + pub fn process_diagnostic( + &self, + input: &DiagnosticInput, + policy: Policy, + ) -> Result { + let message = self.process_text(&input.message, policy)?; + let exception = input + .exception + .as_deref() + .map(|text| self.process_text(text, policy)) + .transpose()?; + let stack = input + .stack + .as_deref() + .map(|text| { + if policy.redact { + self.redact_text(text) + } else { + Ok(text.to_owned()) + } + }) + .transpose()?; + let leaves = input + .leaves + .iter() + .map(|(key, text)| { + if policy.redact { + self.redact_structured_text(key.as_deref(), text) + } else { + Ok(text.clone()) + } + }) + .collect::>>()?; + let changed = message != input.message + || exception != input.exception + || stack != input.stack + || leaves + .iter() + .zip(&input.leaves) + .any(|(processed, (_, original))| processed != original); + Ok(DiagnosticOutput { + message, + exception, + stack, + leaves, + changed, + }) + } + + pub fn scrub_access_arguments(&self, arguments: &[String]) -> Result> { + arguments + .iter() + .map(|argument| self.scrub_access_arg(argument)) + .collect() + } + + fn process_text(&self, text: &str, policy: Policy) -> Result { + let collapsed = if policy.base64_limit > 0 { + collapse_base64(text, policy.base64_limit as usize) + } else { + text.to_owned() + }; + let redacted = if policy.redact { + self.redact_text(&collapsed)? + } else { + collapsed + }; + Ok( + if policy.text_limit > 0 && redacted.chars().count() > policy.text_limit as usize { + truncate_text(&redacted, policy.text_limit as usize) + } else { + redacted + }, + ) + } + + fn scrub_access_arg(&self, value: &str) -> Result { + let length = value.chars().count(); + let scanned = if length <= 512 { + value + } else { + let head = &value[..char_offset(value, 512)]; + if head.contains('?') { + &head[..head.rfind(['?', '&']).unwrap_or(0)] + } else { + head + } + }; + let scrubbed = self.redact_text(scanned)?; + let (path, query) = scrubbed + .split_once('?') + .map_or((scrubbed.as_str(), None), |(path, query)| { + (path, Some(query)) + }); + let safe = if self.hides_encoded_credential(path)? { + REDACTED.to_owned() + } else if query.is_some() && self.hides_encoded_credential(&scrubbed)? { + format!("{path}?{REDACTED}") + } else { + scrubbed + }; + Ok(if length > 512 { + format!( + "{safe}... ({} more chars truncated) ...", + length - scanned.chars().count() + ) + } else { + safe + }) + } + + fn hides_encoded_credential(&self, value: &str) -> Result { + if !value.as_bytes().contains(&b'%') { + return Ok(false); + } + let decoded = percent_decode_str(value).decode_utf8_lossy(); + Ok(self.redact_text(&decoded)? != decoded) + } +} + +fn char_offset(text: &str, count: usize) -> usize { + text.char_indices() + .nth(count) + .map_or(text.len(), |(index, _)| index) +} + +fn marker(skipped_chars: usize) -> String { + format!( + "... (litellm_truncated skipped {skipped_chars} chars. Truncation is a stdout logging safeguard. Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.) and at DEBUG level. To increase the truncation limit, set `MAX_STRING_LENGTH_STDOUT_LOG` in your env.) ..." + ) +} + +fn truncate_text(text: &str, limit: usize) -> String { + let length = text.chars().count(); + let kept = limit.saturating_sub(marker(length).len()); + if kept == 0 { + return text[..char_offset(text, limit)].to_owned(); + } + let head = kept / 2; + let tail = kept - head; + format!( + "{}{}{}", + &text[..char_offset(text, head)], + marker(length - kept), + &text[char_offset(text, length - tail)..] + ) +} + +fn base64_byte(byte: u8) -> bool { + byte.is_ascii_alphanumeric() || byte == b'+' || byte == b'/' +} + +fn looks_like_base64(run: &str) -> bool { + let unpadded = run.trim_end_matches('='); + let lower_hex = unpadded + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte)); + let upper_hex = unpadded + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte)); + let repeated = unpadded.bytes().all(|byte| byte == unpadded.as_bytes()[0]); + (!lower_hex && !upper_hex) || repeated +} + +fn base64_size(chars: usize) -> String { + let bytes = chars as f64 * 3.0 / 4.0; + if bytes >= 1024.0 * 1024.0 { + return format!("{:.2}MB", bytes / (1024.0 * 1024.0)); + } + if bytes >= 1024.0 { + return format!("{:.1}KB", bytes / 1024.0); + } + format!("{}B", bytes as usize) +} + +fn collapse_base64(text: &str, limit: usize) -> String { + let bytes = text.as_bytes(); + let mut position = 0; + let mut previous = 0; + let mut output = String::new(); + while position < bytes.len() { + if !base64_byte(bytes[position]) || (position > 0 && base64_byte(bytes[position - 1])) { + position += 1; + continue; + } + let start = position; + while position < bytes.len() && base64_byte(bytes[position]) { + position += 1; + } + let run_end = position; + while position < bytes.len() && position - run_end < 2 && bytes[position] == b'=' { + position += 1; + } + let run = &text[start..position]; + if run_end - start > limit && looks_like_base64(run) { + output.push_str(&text[previous..start]); + output.push_str(&format!( + "[base64_data truncated: {}]", + base64_size(run.len()) + )); + previous = position; + } + } + output.push_str(&text[previous..]); + output +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn redaction_precedes_the_text_bound_and_preserves_unicode_character_limits() { + let processor = Processor::new(16); + let secret = format!("sk-{}", "q".repeat(48)); + let text = format!("{}{}{}", "é".repeat(110), secret, "界".repeat(1000)); + let input = DiagnosticInput { + message: text, + exception: None, + stack: None, + leaves: vec![], + }; + let output = processor + .process_diagnostic( + &input, + Policy { + redact: true, + base64_limit: 0, + text_limit: 500, + }, + ) + .unwrap(); + assert!(output.message.chars().count() <= 500); + assert!(!output.message.contains("sk-qq")); + assert!(output.changed); + } + + #[test] + fn base64_collapse_applies_to_debug_and_exceptions_without_touching_hex() { + let processor = Processor::new(16); + let input = DiagnosticInput { + message: format!("image={} digest={}", "Q".repeat(100), "a1".repeat(50)), + exception: Some(format!("upload failed: {}", "Q".repeat(100))), + stack: Some("api_key=secret123".to_owned()), + leaves: vec![(Some("api_key".to_owned()), "secret123".to_owned())], + }; + let output = processor + .process_diagnostic( + &input, + Policy { + redact: true, + base64_limit: 20, + text_limit: 0, + }, + ) + .unwrap(); + assert!(output.message.contains("[base64_data truncated: 75B]")); + assert!(output.message.contains(&"a1".repeat(50))); + assert!( + output + .exception + .unwrap() + .contains("[base64_data truncated: 75B]") + ); + assert_eq!(output.stack.as_deref(), Some(REDACTED)); + assert_eq!(output.leaves, vec![REDACTED]); + } + + #[test] + fn access_arguments_keep_encoded_paths_and_drop_decoded_credentials() { + let processor = Processor::new(16); + let arguments = vec![ + "/v1/models?filter=gpt%2D4o&page=2".to_owned(), + "/v1/models?k%65y=sk%2Dabcdefghijklmnopqrstuvwxyz&page=2".to_owned(), + ]; + assert_eq!( + processor.scrub_access_arguments(&arguments).unwrap(), + vec![arguments[0].clone(), "/v1/models?REDACTED".to_owned()] + ); + } +} diff --git a/litellm-rust/crates/tracing/src/redaction.rs b/litellm-rust/crates/tracing/src/redaction.rs new file mode 100644 index 00000000000..4bef14148d3 --- /dev/null +++ b/litellm-rust/crates/tracing/src/redaction.rs @@ -0,0 +1,140 @@ +use fancy_regex::{NoExpand, Regex}; + +pub const REDACTED: &str = "REDACTED"; + +#[cfg(test)] +const DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH: usize = 16; + +fn secret_patterns(minimum_custom_key_length: usize) -> String { + let sk_suffix_length = minimum_custom_key_length.saturating_sub("sk-".len()); + [ + r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----", + r"\bya29\.[A-Za-z0-9_.~+/-]+", + r#"(?:client_secret|azure_password|azure_username)\s+[^\s,'"})\]{}>]+"#, + r"(?:AKIA|ASIA)[0-9A-Z]{16}", + r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*", + r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", + &format!(r"sk-[A-Za-z0-9\-_]{{{sk_suffix_length},}}"), + r#"(?<=[?&])(?:api[_-]?key|\w*(?:token|password|passwd|client_secret|secret_key|_secret))=[^\s&'"]+"#, + r#"(?:api[_-]?key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]{8,}"#, + r#"(?:x-api-key|api-key)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + r"x-ak-[A-Za-z0-9\-_]{20,}", + r"AIza[0-9A-Za-z\-_]{35}", + r#"(?<=[?&])key=[^\s&'"]{8,}"#, + r#"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + r#"(?<=://)[^\s'":]{0,4096}:[^\s'"]{1,4096}(?=@)"#, + r"dapi[0-9a-f]{32}", + r#"litellm\.[A-Za-z0-9_]*_key['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + r#"private_key['"]?\s*[:=]\s*['"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'"})\]{}>]+)"#, + concat!( + r"(?:master_key|xai_key|database_url|db_url|connection_string|", + r"aws_secret_access_key|aws_session_token|aws_access_key_id|s3_secret_access_key|s3_access_key_id|", + r"signing_key|encryption_key|", + r"auth_token|access_token|refresh_token|", + r"slack_webhook_url|webhook_url|", + r"database_connection_string|", + r"huggingface_token|jwt_secret)", + r#"['"]?\s*[:=]\s*['"]?[^\s,'"})\]{}>]+"#, + ), + r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", + r"(?<=[?&])sig=[A-Za-z0-9%+/=]+", + r#"\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"#, + ] + .join("|") +} + +#[derive(Clone, Debug)] +pub struct SecretRedactor { + pattern: Regex, + internal_pattern: Regex, +} + +impl SecretRedactor { + pub fn new(minimum_custom_key_length: usize) -> Self { + let pattern = Regex::new(&format!( + "(?i){}", + secret_patterns(minimum_custom_key_length) + )) + .expect("secret redaction patterns compile"); + let internal_pattern = Regex::new(concat!( + r#"(?i)/(?:etc|var|opt|usr|home|root|private|Users|tmp|mnt|srv)/[^\s'"\)\]}>,]+|"#, + r#"[A-Za-z]:\\[^\s'"\)\]}>,]+|"#, + r"\b(?:10(?:\.\d{1,3}){3}|172\.(?:1[6-9]|2\d|3[01])(?:\.\d{1,3}){2}|", + r"192\.168(?:\.\d{1,3}){2}|127(?:\.\d{1,3}){3})\b|", + r"\b[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.(?:internal|local|corp|lan|intra|private)\b", + )) + .expect("internal detail patterns compile"); + Self { + pattern, + internal_pattern, + } + } + + pub fn redact(&self, value: &str) -> String { + self.try_redact(value) + .unwrap_or_else(|_| REDACTED.to_owned()) + } + + pub fn try_redact(&self, value: &str) -> fancy_regex::Result { + self.pattern + .try_replacen(value, 0, NoExpand(REDACTED)) + .map(|value| value.into_owned()) + } + + pub fn try_redact_structured( + &self, + key: Option<&str>, + value: &str, + ) -> fancy_regex::Result { + let scrubbed = self.try_redact(value)?; + if scrubbed != value || key.is_none() { + return Ok(scrubbed); + } + let rendered = format!("'{}': '{value}'", key.unwrap_or_default()); + Ok(if self.try_redact(&rendered)? != rendered { + REDACTED.to_owned() + } else { + value.to_owned() + }) + } + + pub fn try_redact_internal(&self, value: &str) -> fancy_regex::Result { + let without_traceback = value + .split_once("Traceback (most recent call last):") + .map_or(value, |(prefix, _)| prefix.trim_end()); + self.internal_pattern + .try_replacen(&self.try_redact(without_traceback)?, 0, NoExpand(REDACTED)) + .map(|value| value.into_owned()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[rstest::rstest] + #[case::bearer("auth failed: Bearer abcdefghijklmnop", "auth failed: REDACTED")] + #[case::sk_key("key sk-abcdefghijklmnopqrstuvwxyz rejected", "key REDACTED rejected")] + #[case::short_sk_key_is_kept("sk-abc", "sk-abc")] + #[case::query_param("GET /v1?api_key=secret123&x=1", "GET /v1?REDACTED&x=1")] + #[case::dict_repr("{'api_key': 'abcdefghij'}", "{'REDACTED'}")] + #[case::url_credentials("postgres://user:pass@host/db", "postgres://REDACTED@host/db")] + #[case::case_insensitive("BEARER ABCDEFGHIJKLMNOP", "REDACTED")] + #[case::aws_key("AKIAABCDEFGHIJKLMNOP", "REDACTED")] + #[case::sas_signature("https://x.blob/a?sv=1&sig=abc%2B=", "https://x.blob/a?sv=1&REDACTED")] + #[case::password_needs_word_boundary("db_password=hunter2", "REDACTED")] + #[case::plain_text_is_kept(r#"{"message": "rejected"}"#, r#"{"message": "rejected"}"#)] + fn redacts_the_same_spans_as_the_python_patterns(#[case] input: &str, #[case] expected: &str) { + assert_eq!( + SecretRedactor::new(DEFAULT_MINIMUM_CUSTOM_KEY_LENGTH).redact(input), + expected + ); + } + + #[test] + fn sk_threshold_follows_the_minimum_custom_key_length() { + let redactor = SecretRedactor::new(8); + assert_eq!(redactor.redact("sk-abcde"), REDACTED); + assert_eq!(redactor.redact("sk-abcd"), "sk-abcd"); + } +} diff --git a/litellm-rust/crates/tracing/tests/logging.rs b/litellm-rust/crates/tracing/tests/logging.rs new file mode 100644 index 00000000000..585e442dad1 --- /dev/null +++ b/litellm-rust/crates/tracing/tests/logging.rs @@ -0,0 +1,122 @@ +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + mpsc, +}; + +use litellm_tracing::{Level, Logger, Metadata, Record, Sink, info, warn}; +use serde_json::{Value, json}; + +struct Output { + enabled: Arc, + sender: mpsc::Sender<(String, Value, Level, &'static str, Option)>, +} + +impl Sink for Output { + fn enabled(&self, _: &Metadata<'_>) -> bool { + self.enabled.load(Ordering::Relaxed) + } + + fn emit(&self, record: &Record) { + self.sender + .send(( + record.message.clone(), + Value::Object(record.fields.clone()), + *record.metadata.level(), + record.metadata.target(), + record.metadata.line(), + )) + .unwrap(); + Logger::default().scope(|| warn!("a sink must not recursively emit")); + } +} + +fn emit() { + warn!( + attempt = 3_u64, + elapsed = 1.5, + retry = true, + reason = "timeout", + "retry {}", + 3 + ); +} + +#[test] +fn records_preserve_fields_metadata_and_dynamic_filtering_without_recursion() { + let (sender, receiver) = mpsc::channel(); + let enabled = Arc::new(AtomicBool::new(false)); + let logger = Logger::new(Output { + enabled: enabled.clone(), + sender, + }); + logger.scope(emit); + assert!(receiver.try_recv().is_err()); + enabled.store(true, Ordering::Relaxed); + logger.scope(emit); + let (message, fields, level, target, line) = receiver.try_recv().unwrap(); + assert_eq!(message, "retry 3"); + assert_eq!( + fields, + json!({"attempt": 3, "elapsed": 1.5, "retry": true, "reason": "timeout"}) + ); + assert_eq!(level, Level::WARN); + assert_eq!(target, module_path!()); + assert!(line.is_some()); + enabled.store(false, Ordering::Relaxed); + logger.scope(emit); + assert!(receiver.try_recv().is_err()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn concurrent_futures_keep_their_sinks_across_suspension_and_spawn() { + let tasks = (0..2) + .map(|id| { + let (sender, receiver) = mpsc::channel(); + let logger = Logger::new(Output { + enabled: Arc::new(AtomicBool::new(true)), + sender, + }); + let task = tokio::spawn(logger.instrument(async move { + tokio::task::yield_now().await; + info!(id, "worker"); + })); + (id, task, receiver) + }) + .collect::>(); + for (id, task, receiver) in tasks { + task.await.unwrap(); + let (message, fields, level, _, _) = receiver.try_recv().unwrap(); + assert_eq!(message, "worker"); + assert_eq!(fields, json!({"id": id})); + assert_eq!(level, Level::INFO); + assert!(receiver.try_recv().is_err()); + } +} + +#[test] +fn nested_scopes_restore_the_previous_sink() { + let (outer_sender, outer) = mpsc::channel(); + let (inner_sender, inner) = mpsc::channel(); + let logger = |sender| { + Logger::new(Output { + enabled: Arc::new(AtomicBool::new(true)), + sender, + }) + }; + let outside = logger(outer_sender); + let inside = logger(inner_sender); + outside.scope(|| { + info!("before"); + inside.scope(|| info!("inside")); + info!("after"); + }); + assert_eq!( + outer.try_iter().map(|event| event.0).collect::>(), + ["before", "after"] + ); + assert_eq!( + inner.try_iter().map(|event| event.0).collect::>(), + ["inside"] + ); +} diff --git a/litellm/_logging.py b/litellm/_logging.py index 644a79d8cbd..8bd7a86ecf5 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,11 +1,12 @@ import ast import contextvars import functools +import itertools import logging import os import re import sys -from collections.abc import Sequence +from collections.abc import Iterator from datetime import datetime from logging import Formatter from typing import Any, Final, TextIO @@ -22,10 +23,12 @@ from litellm.litellm_core_utils.env_utils import get_env_int from litellm.litellm_core_utils.safe_json_dumps import UNSERIALIZABLE_OBJECT, safe_dumps, safe_json_structure from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.secret_redaction import ( + _python_redact_string, + _python_redact_structured_value, redact_internal_details, redact_string, - redact_structured_value, ) +from litellm.rust_bridge import diagnostics set_verbose = False @@ -49,7 +52,7 @@ def _sanitize_correlation_id(value: str) -> str: pass through credential redaction. """ stripped: Final = "".join(ch for ch in value if ch.isprintable()) - return _redact_string(stripped[:_MAX_CORRELATION_ID_LENGTH]) + return _redact_string(stripped)[:_MAX_CORRELATION_ID_LENGTH] def set_session_id(session_id: str) -> "contextvars.Token[str]": @@ -74,12 +77,6 @@ def _redact_string(value: str) -> str: return redact_string(value) -def _redact_structured_value(key: str | None, value: str) -> str: - if not _ENABLE_SECRET_REDACTION: - return value - return redact_structured_value(key, value) - - _REDACTED_RECORD_ATTR: Final = "litellm_redacted" _REDACTED_STAMP: Final = object() _UNREDACTED_SCALAR_TYPES: Final = (bool, int, float, type(None)) @@ -103,14 +100,6 @@ def _plain_text(value: object) -> str: return UNSERIALIZABLE_OBJECT -def _redact_extra_value(key: str, value: object) -> object: - try: - scrubbed: Final = safe_json_structure(value, value_transform=_redact_structured_value, key=key) - except Exception: - return _redact_string(_plain_text(value)) - return value if _scrubbing_changed_nothing(scrubbed, value) else scrubbed - - def redact_secrets(value: str) -> str: """Public API: redact known secret/credential patterns from an arbitrary string. @@ -150,7 +139,7 @@ def _substituted_color_message(record: logging.LogRecord) -> str | None: return None try: return color_message % record.args - except TypeError: + except Exception: return color_message @@ -162,42 +151,7 @@ class SecretRedactionFilter(logging.Filter): def filter(self, record: logging.LogRecord) -> bool: if not _ENABLE_SECRET_REDACTION or _is_redacted(record): return True - - # Runs before args are cleared, and before the extra-field loop below - # that redacts the substituted result. - substituted_color_message: Final = _substituted_color_message(record) - if substituted_color_message is not None: - record.color_message = substituted_color_message # rebind-ok: a Filter scrubs records in place - - try: - record.msg = _redact_string(record.getMessage()) - record.args = None - except Exception: - if isinstance(record.msg, str): - record.msg = _redact_string(record.msg) - - # Redact exception tracebacks - if record.exc_info and record.exc_info[1] is not None: - try: - record.exc_text = _redact_string(record.exc_text or self._formatter.formatException(record.exc_info)) - except Exception: - pass - - if isinstance(record.stack_info, str): - record.stack_info = _redact_string(record.stack_info) # rebind-ok: a Filter scrubs records in place - - # Redact extra fields passed via logger.debug("msg", extra={...}) - record_items: Final[Sequence[tuple[str, object]]] = list(record.__dict__.items()) - for key, value in record_items: - if key in _STANDARD_RECORD_ATTRS: - continue - if isinstance(value, str): - setattr(record, key, _redact_structured_value(key, value)) - elif not isinstance(value, _UNREDACTED_SCALAR_TYPES): - setattr(record, key, _redact_extra_value(key, value)) - - setattr(record, _REDACTED_RECORD_ATTR, _REDACTED_STAMP) - return True + return _process_record(record, base64_limit=0, text_limit=0, redact=True) _secret_filter: Final = SecretRedactionFilter() @@ -211,7 +165,7 @@ _REDACTION_PLACEHOLDER: Final = "REDACTED" def _hides_a_credential(value: str) -> bool: """Whether *value* only looks clean until it is percent-decoded.""" decoded: Final = unquote(value) - return _redact_string(decoded) != decoded + return _python_redact_string(decoded) != decoded def _drop_encoded_credential(scrubbed: str) -> str: @@ -239,10 +193,10 @@ def _scrub_access_arg(value: str) -> str: pattern and would then be logged raw. """ if len(value) <= _MAX_SCRUBBED_ACCESS_ARG: - return _drop_encoded_credential(_redact_string(value)) + return _drop_encoded_credential(_python_redact_string(value)) head: Final = value[:_MAX_SCRUBBED_ACCESS_ARG] kept: Final = head[: max(head.rfind("?"), head.rfind("&"))] if "?" in head else head - scrubbed: Final = _drop_encoded_credential(_redact_string(kept)) + scrubbed: Final = _drop_encoded_credential(_python_redact_string(kept)) return f"{scrubbed}... ({len(value) - len(kept)} more chars truncated) ..." @@ -258,8 +212,17 @@ class AccessLogRedactionFilter(logging.Filter): if not _ENABLE_SECRET_REDACTION: return True if isinstance(record.args, tuple) and record.args: + strings: Final = tuple(arg for arg in record.args if isinstance(arg, str)) + candidate: Final = diagnostics.run( + lambda native: native.scrub_access_arguments(strings), + lambda: tuple(_scrub_access_arg(arg) for arg in strings), + ) + scrubbed: Final = ( + candidate if len(candidate) == len(strings) else tuple(_scrub_access_arg(arg) for arg in strings) + ) + values: Final = iter(scrubbed) record.args = tuple( # rebind-ok: a Filter scrubs records in place - _scrub_access_arg(arg) if isinstance(arg, str) else arg for arg in record.args + next(values) if isinstance(arg, str) else arg for arg in record.args ) return True # No positional args means everything is in msg, where collapsing is correct. @@ -365,6 +328,185 @@ def _collapse_base64_runs(text: str, limit: int) -> str: return _base64_run_pattern(limit + 1).sub(_replace_base64_run, text) +def _extra_structure(key: str, value: object) -> object: + if isinstance(value, str): + return value + try: + return safe_json_structure(value, key=key) + except Exception: + return _plain_text(value) + + +def _string_leaves(key: str | None, value: object) -> Iterator[tuple[str | None, str]]: + if isinstance(value, str): + yield key, value + elif isinstance(value, dict): + yield from itertools.chain.from_iterable( + _string_leaves(child_key, child) for child_key, child in value.items() if isinstance(child_key, str) + ) + elif isinstance(value, (list, tuple)): + yield from itertools.chain.from_iterable(_string_leaves(key, child) for child in value) + + +def _replace_string_leaves(value: object, values: Iterator[str]) -> object: + if isinstance(value, str): + return next(values) + if isinstance(value, dict): + return { # mutable-ok: LogRecord extras must keep JSON dict shape for handlers + key: _replace_string_leaves(child, values) for key, child in value.items() + } + if isinstance(value, list): + return [ # mutable-ok: LogRecord extras must keep JSON list shape for handlers + _replace_string_leaves(child, values) for child in value + ] + if isinstance(value, tuple): + return tuple(_replace_string_leaves(child, values) for child in value) + return value + + +def _sort_processed_sets(original: object, processed: object) -> object: + if isinstance(original, set) and isinstance(processed, list): + return sorted(processed) + if isinstance(original, dict) and isinstance(processed, dict): + return { # mutable-ok: sorting nested sets must preserve the surrounding JSON dict + key: _sort_processed_sets(original.get(key), value) for key, value in processed.items() + } + if isinstance(original, list) and isinstance(processed, list): + return [ # mutable-ok: sorting nested sets must preserve the surrounding JSON list + _sort_processed_sets(before, after) for before, after in zip(original, processed) + ] + if isinstance(original, tuple) and isinstance(processed, tuple): + return tuple(_sort_processed_sets(before, after) for before, after in zip(original, processed)) + return processed + + +def _python_process_diagnostic( + message: str, + exception: str | None, + stack: str | None, + leaves: tuple[tuple[str | None, str], ...], + redact: bool, + base64_limit: int, + text_limit: int, +) -> tuple[str, str | None, str | None, tuple[str, ...], bool]: + def process_text(text: str) -> str: + collapsed: Final = _collapse_base64_runs(text, base64_limit) if base64_limit > 0 else text + scrubbed: Final = _python_redact_string(collapsed) if redact else collapsed + return _truncate_for_stdout_log(scrubbed, text_limit) if 0 < text_limit < len(scrubbed) else scrubbed + + processed_message: Final = process_text(message) + processed_exception: Final = process_text(exception) if exception is not None else None + processed_stack: Final = _python_redact_string(stack) if redact and stack is not None else stack + processed_leaves: Final = tuple( + _python_redact_structured_value(key, text) if redact else text for key, text in leaves + ) + changed: Final = ( + processed_message != message + or processed_exception != exception + or processed_stack != stack + or any(processed != original for processed, (_, original) in zip(processed_leaves, leaves)) + ) + return processed_message, processed_exception, processed_stack, processed_leaves, changed + + +def _render_message(record: logging.LogRecord) -> str: + try: + return record.getMessage() + except Exception: + return record.msg if isinstance(record.msg, str) else UNSERIALIZABLE_OBJECT + + +def _render_exception(record: logging.LogRecord) -> str | None: + if not isinstance(record.exc_info, tuple) or len(record.exc_info) < 2 or record.exc_info[1] is None: + return None + try: + return record.exc_text or SecretRedactionFilter._formatter.formatException(record.exc_info) + except Exception: + return "REDACTED" + + +def _process_record(record: logging.LogRecord, *, base64_limit: int, text_limit: int, redact: bool) -> bool: + if _is_redacted(record): + return True + message: Final = _render_message(record) + exception: Final = _render_exception(record) + stack: Final = record.stack_info if isinstance(record.stack_info, str) else None + substituted_color: Final = _substituted_color_message(record) + extras: Final = ( + tuple( + ( + key, + value, + _extra_structure( + key, substituted_color if key == "color_message" and substituted_color is not None else value + ), + ) + for key, value in record.__dict__.items() + if key not in _STANDARD_RECORD_ATTRS + and key != _REDACTED_RECORD_ATTR + and not isinstance(value, _UNREDACTED_SCALAR_TYPES) + ) + if redact + else () + ) + extra_leaves: Final = tuple( + itertools.chain.from_iterable(_string_leaves(key, prepared) for key, _, prepared in extras) + ) + raw_template: Final = record.msg if redact and isinstance(record.msg, str) and record.args else None + color_template: Final = record.__dict__.get("color_message") + raw_color: Final = color_template if redact and isinstance(color_template, str) and record.args else None + leaves: Final = ( + extra_leaves + + (((None, raw_template),) if raw_template is not None else ()) + + (((None, raw_color),) if raw_color is not None else ()) + ) + candidate: Final = diagnostics.run( + lambda native: native.process_diagnostic(message, exception, stack, leaves, (redact, base64_limit, text_limit)), + lambda: _python_process_diagnostic(message, exception, stack, leaves, redact, base64_limit, text_limit), + ) + processed_message, processed_exception, processed_stack, processed_leaves, _ = ( + candidate + if len(candidate[3]) == len(leaves) + else _python_process_diagnostic(message, exception, stack, leaves, redact, base64_limit, text_limit) + ) + raw_template_changed: Final = raw_template is not None and processed_leaves[len(extra_leaves)] != raw_template + safe_message: Final = "REDACTED" if raw_template_changed and processed_message == message else processed_message + if redact or safe_message != message: + record.msg = safe_message # rebind-ok: the Filter interface mutates the record + record.args = None # rebind-ok: the rendered message replaces interpolation inputs + if processed_exception is not None: + record.exc_text = processed_exception # rebind-ok: the Filter interface mutates the record + if processed_stack is not None: + record.stack_info = processed_stack # rebind-ok: the Filter interface mutates the record + processed_values: Final = iter(processed_leaves[: len(extra_leaves)]) + for key, original, prepared in extras: + replacement: Final = _sort_processed_sets(original, _replace_string_leaves(prepared, processed_values)) + if not _scrubbing_changed_nothing(replacement, original): + setattr(record, key, replacement) + raw_color_changed: Final = ( + raw_color is not None and processed_leaves[len(extra_leaves) + int(raw_template is not None)] != raw_color + ) + if raw_color_changed and getattr(record, "color_message", None) == substituted_color: + setattr(record, "color_message", "REDACTED") + setattr(record, _REDACTED_RECORD_ATTR, _REDACTED_STAMP) + return True + + +def _redact_json_record(value: object) -> object: + prepared: Final = safe_json_structure(value) + leaves: Final = tuple(_string_leaves(None, prepared)) + candidate: Final = diagnostics.run( + lambda native: native.process_diagnostic("", None, None, leaves, (True, 0, 0))[3], + lambda: tuple(_python_redact_structured_value(key, text) for key, text in leaves), + ) + replacements: Final = ( + candidate + if len(candidate) == len(leaves) + else tuple(_python_redact_structured_value(key, text) for key, text in leaves) + ) + return _sort_processed_sets(value, _replace_string_leaves(prepared, iter(replacements))) + + class StdoutLogTruncationFilter(logging.Filter): """Bounds how much of an oversized log line reaches stdout. @@ -412,7 +554,17 @@ class StdoutLogTruncationFilter(logging.Filter): return True -_stdout_truncation_filter: Final = StdoutLogTruncationFilter() +class DiagnosticProcessingFilter(StdoutLogTruncationFilter): + def filter(self, record: logging.LogRecord) -> bool: + return _process_record( + record, + base64_limit=_get_max_base64_length_stdout_log(), + text_limit=_get_max_string_length_stdout_log() if record.levelno >= logging.INFO else 0, + redact=_ENABLE_SECRET_REDACTION, + ) + + +_diagnostic_filter: Final = DiagnosticProcessingFilter() class CorrelationContextFilter(logging.Filter): @@ -633,7 +785,9 @@ class JsonFormatter(Formatter): if record.exc_info: json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info) - return safe_dumps(json_record, value_transform=None if _is_redacted(record) else _redact_structured_value) + return safe_dumps( + json_record if _is_redacted(record) or not _ENABLE_SECRET_REDACTION else _redact_json_record(json_record) + ) class CorrelationPlainFormatter(logging.Formatter): @@ -663,7 +817,7 @@ def _setup_json_exception_handlers(formatter): # Create a handler with JSON formatting for exceptions error_handler: Final = logging.StreamHandler() error_handler.setFormatter(formatter) - error_handler.addFilter(_stdout_truncation_filter) + error_handler.addFilter(_diagnostic_filter) error_handler.addFilter(_secret_filter) error_handler.addFilter(_correlation_filter) @@ -734,10 +888,10 @@ verbose_logger.addHandler(handler) # Filters attached to the logger, not the handler, survive callers swapping in their own # handlers (JSON mode, uvicorn log config, a host app's root handler). -verbose_router_logger.addFilter(_stdout_truncation_filter) -verbose_proxy_logger.addFilter(_stdout_truncation_filter) -verbose_proxy_stdout_logger.addFilter(_stdout_truncation_filter) -verbose_logger.addFilter(_stdout_truncation_filter) +verbose_router_logger.addFilter(_diagnostic_filter) +verbose_proxy_logger.addFilter(_diagnostic_filter) +verbose_proxy_stdout_logger.addFilter(_diagnostic_filter) +verbose_logger.addFilter(_diagnostic_filter) def _suppress_loggers(): diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index 70b2cd08b4c..89df5e500db 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -10,6 +10,7 @@ import re from typing import Final from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH +from litellm.rust_bridge import diagnostics REDACTED: Final = "REDACTED" @@ -87,9 +88,13 @@ def _build_secret_patterns() -> "re.Pattern[str]": _SECRET_RE: Final = _build_secret_patterns() +def _python_redact_string(value: str) -> str: + return _SECRET_RE.sub(REDACTED, value) + + def redact_string(value: str) -> str: """Scrub known secret/credential patterns from *value* and return the result.""" - return _SECRET_RE.sub(REDACTED, value) + return diagnostics.run(lambda native: native.redact_text(value), lambda: _python_redact_string(value)) _UNIX_SYSTEM_PATH: Final = r"/(?:etc|var|opt|usr|home|root|private|Users|tmp|mnt|srv)/[^\s'\"\)\]}>,]+" @@ -105,15 +110,21 @@ _INTERNAL_DETAIL_RE: Final = re.compile( _TRACEBACK_MARKER: Final = "Traceback (most recent call last):" -def redact_internal_details(value: str) -> str: +def _python_redact_internal_details(value: str) -> str: """Drop an embedded traceback and scrub filesystem paths and internal hostnames, on top of redact_string(). For client-facing messages only: server logs keep this detail.""" marker_index: Final = value.find(_TRACEBACK_MARKER) without_traceback: Final = value[:marker_index].rstrip() if marker_index != -1 else value - return _INTERNAL_DETAIL_RE.sub(REDACTED, redact_string(without_traceback)) + return _INTERNAL_DETAIL_RE.sub(REDACTED, _python_redact_string(without_traceback)) -def redact_structured_value(key: str | None, value: str) -> str: +def redact_internal_details(value: str) -> str: + return diagnostics.run( + lambda native: native.redact_client_message(value), lambda: _python_redact_internal_details(value) + ) + + +def _python_redact_structured_value(key: str | None, value: str) -> str: """Scrub *value* as it appeared under *key* inside a structured record. redact_string() replaces a whole ``key: value`` span with REDACTED, which is @@ -122,8 +133,15 @@ def redact_structured_value(key: str | None, value: str) -> str: repr would, so the key-name patterns still fire, but collapses only the value so the caller's structure survives. """ - scrubbed: Final = redact_string(value) + scrubbed: Final = _python_redact_string(value) if scrubbed != value or key is None: return scrubbed rendered: Final = f"'{key}': '{value}'" - return REDACTED if redact_string(rendered) != rendered else value + return REDACTED if _python_redact_string(rendered) != rendered else value + + +def redact_structured_value(key: str | None, value: str) -> str: + return diagnostics.run( + lambda native: native.redact_structured_text(key, value), + lambda: _python_redact_structured_value(key, value), + ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index f794dc4a99e..355a5e5062e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -12,6 +12,22 @@ class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... +@final +class NativeDiagnosticProcessor: + def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ... + def redact_text(self, text: str) -> str: ... + def redact_structured_text(self, key: str | None, text: str) -> str: ... + def redact_client_message(self, text: str) -> str: ... + def process_diagnostic( + self, + message: str, + exception: str | None, + stack: str | None, + leaves: Sequence[tuple[str | None, str]], + policy: tuple[bool, int, int], + ) -> tuple[str, str | None, str | None, list[str], bool]: ... + def scrub_access_arguments(self, arguments: Sequence[str]) -> list[str]: ... + def ocr( request: LiteLLMOcrRequest, args: tuple[object, ...], @@ -335,6 +351,7 @@ def reserve_process_for_forking() -> None: ... __all__ = [ "ForkedAfterNativeRuntimeStarted", "HuggingFaceEncoding", + "NativeDiagnosticProcessor", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index d6cbe17d790..1a2d153871d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -86,11 +86,25 @@ class SecretManagerRule: return isinstance(context, SecretManagerContext) and (self.systems is None or context.system in self.systems) -Context: TypeAlias = RouteContext | CacheContext | SecretManagerContext -Rule: TypeAlias = RouteRule | CacheRule | SecretManagerRule +@dataclass(frozen=True, slots=True) +class LoggerContext: + pass + + +@dataclass(frozen=True, slots=True) +class LoggerRule: + rollout: Rollout + + def matches(self, context: Context) -> bool: + return isinstance(context, LoggerContext) + + +Context: TypeAlias = RouteContext | CacheContext | SecretManagerContext | LoggerContext +Rule: TypeAlias = RouteRule | CacheRule | SecretManagerRule | LoggerRule Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( + LoggerRule(Rollout.RUST_OPT_IN), RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})), RouteRule(Route.OCR, Rollout.RUST_OPT_OUT), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), diff --git a/litellm/rust_bridge/diagnostics.py b/litellm/rust_bridge/diagnostics.py new file mode 100644 index 00000000000..fbe1b8122d7 --- /dev/null +++ b/litellm/rust_bridge/diagnostics.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from functools import lru_cache +from typing import Final, Protocol, TypeVar, cast + +from litellm.rust_bridge.bindings import NativeBinding + +ResultT: Final = TypeVar("ResultT") + + +class NativeDiagnosticProcessor(Protocol): + def redact_text(self, text: str) -> str: ... + def redact_structured_text(self, key: str | None, text: str) -> str: ... + def redact_client_message(self, text: str) -> str: ... + def process_diagnostic( + self, + message: str, + exception: str | None, + stack: str | None, + leaves: tuple[tuple[str | None, str], ...], + policy: tuple[bool, int, int], + ) -> tuple[str, str | None, str | None, Sequence[str], bool]: ... + def scrub_access_arguments(self, arguments: tuple[str, ...]) -> Sequence[str]: ... + + +class NativeDiagnosticFactory(Protocol): + def __call__(self, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ... + + +def _as_factory(value: object) -> NativeDiagnosticFactory | None: + if not isinstance(value, type): + return None + return cast(NativeDiagnosticFactory, value) # cast-ok: PyO3 factory must be a type + + +PROCESSOR: Final = NativeBinding("NativeDiagnosticProcessor", validate=_as_factory) + + +@lru_cache(maxsize=4) +def _construct(factory: NativeDiagnosticFactory, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: + return factory(minimum_custom_key_length) + + +def run(native: Callable[[NativeDiagnosticProcessor], ResultT], python: Callable[[], ResultT]) -> ResultT: + from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH + from litellm.rust_bridge.catalog import LoggerContext, decision + from litellm.rust_bridge.configuration import Decision + + selected: Final = decision(LoggerContext()) + if selected is Decision.PYTHON: + return python() + factory: Final = PROCESSOR.load() + if factory is None: + return python() + try: + return native(_construct(factory, MINIMUM_CUSTOM_KEY_LENGTH)) + except Exception: + return python() diff --git a/litellm/rust_bridge/logger.py b/litellm/rust_bridge/logger.py new file mode 100644 index 00000000000..bcd53852a2f --- /dev/null +++ b/litellm/rust_bridge/logger.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +from pydantic import JsonValue + +from litellm._logging import ( + CorrelationContextFilter, + DiagnosticProcessingFilter, + session_id_var, + set_session_id, + set_trace_id, + trace_id_var, + verbose_logger, +) + +_REDACTION: Final = DiagnosticProcessingFilter() +_CORRELATION: Final = CorrelationContextFilter() + + +def context() -> tuple[str, str]: + return session_id_var.get(), trace_id_var.get() + + +def enabled(level: int) -> bool: + return verbose_logger.isEnabledFor(level) + + +def emit( + level: int, + message: str, + pathname: str, + lineno: int, + target: str, + fields: Mapping[str, JsonValue], + correlation: tuple[str, str], +) -> None: + if not enabled(level): + return + session_token: Final = set_session_id(correlation[0]) + trace_token: Final = set_trace_id(correlation[1]) + try: + record: Final = verbose_logger.makeRecord( + verbose_logger.name, + level, + pathname, + lineno, + message, + (), + None, + func=target, + extra={ + "rust_target": target, + "rust_fields": dict(fields), + }, # mutable-ok: LogRecord requires JSON dict extras + ) + _REDACTION.filter(record) + _CORRELATION.filter(record) + verbose_logger.handle(record) + finally: + trace_id_var.reset(trace_token) + session_id_var.reset(session_token) diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 866f2fce989..46e96b60620 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -58,12 +58,6 @@ class SecretManagerBinding: settings_object: object -def warn(message: str) -> None: - from litellm._logging import verbose_logger - - verbose_logger.warning("%s", message) - - def secret_manager() -> SecretManager: from litellm.secret_managers.main import ( _should_read_secret_from_secret_manager, # pyright: ignore[reportPrivateUsage] # canonical resolver is private diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index b77348b5d73..dc1f8592612 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -69,6 +69,9 @@ IGNORE_FUNCTIONS = [ "_unqualified", # bounded by the qualifier depth of a static TypedDict annotation (Annotated, Required/NotRequired, ReadOnly around one type, no cycles possible). "_render_json", # bounded by the nesting depth of a pydantic-validated JsonValue from the operator's config (a finite JSON tree, no cycles possible). "completion_cost", # max depth 1: recursion only fires for mixed-tier Responses WS logging objects, and each split part carries a single service_tier so _split_responses_ws_logging_object_by_service_tier returns None. + "_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). + "_replace_string_leaves", # bounded by the nesting depth of a safe_json_structure output (a finite JSON tree, no cycles possible). + "_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible). ] diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index b6b486e73f0..45c35dc6802 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -11,6 +11,7 @@ from litellm.rust_bridge.catalog import ( CacheRule, Context, Delivery, + LoggerContext, Route, RouteContext, RouteRule, @@ -89,6 +90,13 @@ def test_backend_rollouts_stay_on_python_when_global_rust_is_enabled( assert catalog.decision(context) is Decision.PYTHON +def test_logger_rollout_obeys_the_global_switch() -> None: + assert catalog.rollout(LoggerContext()) is Rollout.RUST_OPT_IN + assert catalog.decision(LoggerContext()) is Decision.PYTHON + configuration.rust(True) + assert catalog.decision(LoggerContext()) is Decision.RUST_WITH_FALLBACK + + def test_response_cache_rules_select_the_whole_backend_runtime() -> None: rules: Final = ( CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), diff --git a/tests/test_litellm/rust_bridge/test_logger.py b/tests/test_litellm/rust_bridge/test_logger.py new file mode 100644 index 00000000000..aff385b9344 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_logger.py @@ -0,0 +1,137 @@ +import logging +from typing import Final + +import pytest + +import litellm +from litellm._logging import ( + DiagnosticProcessingFilter, + _python_process_diagnostic, + redact_secrets, + session_id_var, + trace_id_var, + verbose_logger, +) +from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH +from litellm.litellm_core_utils.secret_redaction import ( + _python_redact_internal_details, + _python_redact_string, + _python_redact_structured_value, +) +from litellm.rust_bridge import diagnostics, logger + + +def test_native_records_preserve_metadata_and_redact_before_custom_handlers( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(verbose_logger, "handlers", []) + secret: Final = "sk-" + "a" * 48 + message: Final = f"Authorization: Bearer {secret}" + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logger.emit( + logging.WARNING, message, "native.rs", 42, "litellm_http", {"retry": True, "api_key": secret}, ("", "") + ) + + record: Final = caplog.records[0] + assert len(caplog.records) == 1 + assert record.getMessage() == redact_secrets(message) + assert secret not in record.getMessage() + assert (record.pathname, record.lineno, record.funcName) == ("native.rs", 42, "litellm_http") + assert record.__dict__["rust_fields"]["retry"] is True + assert secret not in str(record.__dict__["rust_fields"]) + + +def test_native_context_is_scoped_and_respects_correlation_setting( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + context_before: Final = logger.context() + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + logger.emit(logging.WARNING, "native", "native.rs", 1, "litellm_http", {}, ("session", "trace")) + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + logger.emit(logging.WARNING, "disabled", "native.rs", 2, "litellm_http", {}, ("hidden", "hidden")) + + first, second = caplog.records + assert (first.__dict__["session_id"], first.__dict__["trace_id"]) == ("session", "trace") + assert "session_id" not in second.__dict__ + assert "trace_id" not in second.__dict__ + assert (session_id_var.get(), trace_id_var.get()) == context_before + + +def test_native_logging_observes_level_changes(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.ERROR, logger="LiteLLM"): + assert not logger.enabled(logging.WARNING) + logger.emit(logging.WARNING, "filtered", "native.rs", 1, "litellm_http", {}, ("", "")) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert logger.enabled(logging.WARNING) + logger.emit(logging.WARNING, "visible", "native.rs", 1, "litellm_http", {}, ("", "")) + + assert [record.getMessage() for record in caplog.records] == ["visible"] + + +@pytest.mark.parametrize( + "text", + ( + "Authorization: Bearer abcdefghijklmnop", + "s3_secret_access_key=secret123", + "postgres://user:pass@database.internal/name", + '{"type":"service_account","private_key":"secret123"}', + "GET /v1?key=abcdefghij&page=2", + ), +) +def test_native_credential_patterns_match_python(text: str) -> None: + pytest.importorskip("litellm.rust_bridge._native") + from litellm.rust_bridge._native import NativeDiagnosticProcessor + + processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH) + assert processor.redact_text(text) == _python_redact_string(text) + assert processor.redact_structured_text("api_key", "secret123") == _python_redact_structured_value( + "api_key", "secret123" + ) + + +def test_native_client_redaction_matches_python() -> None: + pytest.importorskip("litellm.rust_bridge._native") + from litellm.rust_bridge._native import NativeDiagnosticProcessor + + text: Final = "error at /etc/secrets/config on db.internal\nTraceback (most recent call last):\nsecret" + processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH) + assert processor.redact_client_message(text) == _python_redact_internal_details(text) + + +def test_native_diagnostic_batch_matches_python() -> None: + pytest.importorskip("litellm.rust_bridge._native") + from litellm.rust_bridge._native import NativeDiagnosticProcessor + + message: Final = "é" * 110 + "sk-" + "q" * 48 + "界" * 1000 + exception: Final = "document=" + "Q" * 200 + stack: Final = "api_key=secret123" + leaves: Final = (("api_key", "secret123"), (None, "safe")) + processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH) + rust: Final = processor.process_diagnostic(message, exception, stack, leaves, (True, 20, 500)) + python: Final = _python_process_diagnostic(message, exception, stack, leaves, True, 20, 500) + + assert rust[:3] == python[:3] + assert tuple(rust[3]) == python[3] + assert rust[4] == python[4] + assert "sk-qq" not in rust[0] + assert len(rust[0]) <= 500 + assert rust[3] == ["REDACTED", "safe"] + + +def test_missing_native_diagnostic_processor_falls_back_before_record_mutation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + diagnostics.PROCESSOR.override(None) + try: + record: Final = logging.makeLogRecord({"name": "LiteLLM", "levelno": logging.INFO, "msg": "api_key=secret123"}) + assert DiagnosticProcessingFilter().filter(record) is True + assert record.getMessage() == "REDACTED" + finally: + diagnostics.PROCESSOR.reset() + + +def test_unsupported_unicode_uses_safe_python_redaction(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + assert redact_secrets("broken\ud800 api_key=secret123") == "broken\ud800 REDACTED" diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 3a86a69ed8b..216961ab640 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -1,4 +1,3 @@ -import logging from typing import Final import httpx @@ -11,6 +10,7 @@ from litellm.rust_bridge import settings from litellm.secret_managers.main import get_secret_str from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem + def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "user_url_validation", False) monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"]) @@ -57,13 +57,6 @@ def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyP assert result.ssl_verify is True -def test_warn_reaches_the_litellm_logger(caplog: pytest.LogCaptureFixture) -> None: - with caplog.at_level(logging.WARNING, logger="LiteLLM"): - settings.warn("ssl_ecdh_curve 'secp521r1' is not supported") - - assert [record.getMessage() for record in caplog.records] == ["ssl_ecdh_curve 'secp521r1' is not supported"] - - class _VaultSecrets(CustomSecretManager): def __init__(self, secrets: dict[str, str]) -> None: super().__init__(secret_manager_name="rust_bridge_settings_test") diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 7cecdaec25d..d9cfe88d52f 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -19,6 +19,16 @@ from litellm._logging import ( _COLOR_LOG_FORMAT, _MAX_SCRUBBED_ACCESS_ARG, _PLAIN_LOG_FORMAT, + ALL_LOGGERS, + AccessLogPathFilter, + AccessLogRedactionFilter, + CorrelationContextFilter, + CorrelationPlainFormatter, + DiagnosticProcessingFilter, + JsonFormatter, + LevelRoutingStreamHandler, + SecretRedactionFilter, + StdoutLogTruncationFilter, _get_uvicorn_json_log_config, _initialize_loggers_with_handler, _parse_json_logs_env, @@ -33,15 +43,6 @@ from litellm._logging import ( verbose_logger, verbose_proxy_logger, verbose_router_logger, - ALL_LOGGERS, - AccessLogPathFilter, - AccessLogRedactionFilter, - CorrelationContextFilter, - CorrelationPlainFormatter, - JsonFormatter, - LevelRoutingStreamHandler, - SecretRedactionFilter, - StdoutLogTruncationFilter, ) from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD from litellm.integrations.custom_logger import CustomLogger @@ -824,6 +825,75 @@ def test_secret_filter_keeps_truncated_traceback(monkeypatch): assert "sk-1234567890abcdefghij" not in record.exc_text +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_diagnostic_redaction_precedes_a_credential_cut(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500") + monkeypatch.setenv("MAX_BASE64_LENGTH_STDOUT_LOG", "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + secret = "sk-" + "q" * 48 + record = _make_record(logging.INFO, "%s", ("é" * 110 + secret + "界" * 1000,)) + + assert DiagnosticProcessingFilter().filter(record) is True + + assert len(record.getMessage()) <= 500 + assert "sk-qq" not in record.getMessage() + + +def test_correlation_id_redacts_before_its_length_bound(monkeypatch): + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + secret = "sk-" + "q" * 48 + token = set_trace_id("x" * 250 + secret) + try: + assert "sk-qq" not in trace_id_var.get() + assert len(trace_id_var.get()) <= 256 + finally: + trace_id_var.reset(token) + + +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + record = _make_record(logging.WARNING, "bad % api_key=secret123", ("value",)) + record.color_message = "bad % api_key=secret123" + + assert DiagnosticProcessingFilter().filter(record) is True + + assert record.getMessage() == "REDACTED" + assert record.color_message == "REDACTED" + + +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_key_pattern_template_keeps_the_rendered_redacted_line(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + record = _make_record(logging.INFO, "password=%s ok", ("hunter2",)) + record.color_message = "password=%s ok" + + assert DiagnosticProcessingFilter().filter(record) is True + + assert record.getMessage() == "REDACTED ok" + assert record.color_message == "REDACTED ok" + + +def test_disabled_diagnostic_call_does_not_render_arguments(caplog): + class Unrenderable: + def __str__(self): + raise AssertionError("disabled call rendered its argument") + + with caplog.at_level(logging.ERROR, logger="LiteLLM"): + verbose_logger.debug("hidden %s", Unrenderable()) + + assert not caplog.records + + def test_truncation_filter_survives_json_reconfiguration(): """The cap lives on the loggers, so swapping handlers (JSON mode) can't drop it.""" _turn_on_json() @@ -983,10 +1053,10 @@ _REQUEST_DUMP = "{'model': 'gpt-4', 'messages': [{'role': 'user', 'content': 'he (CorrelationPlainFormatter(_PLAIN_LOG_FORMAT), JsonFormatter()), ids=("plain", "json"), ) -def test_scrubbed_record_is_scanned_for_secrets_once(monkeypatch, formatter): - """Every pass of the secret regex over a multi-megabyte debug line costs seconds of - event-loop time, so a formatter must not rescan what SecretRedactionFilter scrubbed.""" +def test_scrubbed_record_scans_the_large_rendered_value_once(monkeypatch, formatter): + """The raw format template gets its own check, while the large rendered value gets one scan.""" counting = _CountingPattern(secret_redaction._SECRET_RE) + monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setattr(secret_redaction, "_SECRET_RE", counting) monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.DEBUG, "receiving data: %s", (_REQUEST_DUMP,)) @@ -997,14 +1067,15 @@ def test_scrubbed_record_is_scanned_for_secrets_once(monkeypatch, formatter): assert _REQUEST_DUMP in rendered assert "litellm_redacted" not in rendered - assert counting.calls == 1 - assert counting.scanned_chars == len(f"receiving data: {_REQUEST_DUMP}") + assert counting.calls == 2 + assert counting.scanned_chars == len(f"receiving data: {_REQUEST_DUMP}") + len("receiving data: %s") def test_stamped_record_is_not_scanned_again(monkeypatch): """JSON mode puts the filter on a third-party logger and again on the root handler its records propagate to, so the second filter must trust the stamp instead of rescanning.""" counting = _CountingPattern(secret_redaction._SECRET_RE) + monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setattr(secret_redaction, "_SECRET_RE", counting) monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.DEBUG, "receiving data: %s", (_REQUEST_DUMP,)) @@ -1012,13 +1083,14 @@ def test_stamped_record_is_not_scanned_again(monkeypatch): assert SecretRedactionFilter().filter(record) is True assert SecretRedactionFilter().filter(record) is True - assert counting.calls == 1 + assert counting.calls == 2 def test_caller_supplied_stamp_never_skips_the_scrub(monkeypatch): """The stamp is a private sentinel, so a caller passing extra={"litellm_redacted": True} still gets the full scrub, and only the filter's own stamp lets a later pass skip it.""" counting = _CountingPattern(secret_redaction._SECRET_RE) + monkeypatch.setenv("LITELLM_RUST", "0") monkeypatch.setattr(secret_redaction, "_SECRET_RE", counting) monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.DEBUG, "api_key=sk-1234567890abcdefghij") @@ -1619,3 +1691,75 @@ def test_access_log_path_filter_keeps_a_record_without_a_string_path_arg(monkeyp exc_info=None, ) assert AccessLogPathFilter().filter(record) is True + + +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + secret = "sk-" + "q" * 48 + try: + raise ValueError(f"upstream rejected {secret}") + except ValueError: + record = _make_record(logging.ERROR, "call failed", exc_info=sys.exc_info()) + record.stack_info = f"Stack (most recent call last): {secret}" + record.payload = { + "api_key": secret, + "items": [secret, "ok"], + "tags": {secret}, + "pair": (secret, "ok"), + "count": 2, + } + + assert DiagnosticProcessingFilter().filter(record) is True + + assert secret not in (record.exc_text or "") + assert secret not in (record.stack_info or "") + assert secret not in repr(record.payload) + assert record.payload["count"] == 2 + + +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + record = _make_record(logging.WARNING, "api_key=secret123") + diagnostic_filter = DiagnosticProcessingFilter() + + assert diagnostic_filter.filter(record) is True + assert diagnostic_filter.filter(record) is True + assert record.getMessage() == "REDACTED" + + +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + secret = "sk-" + "q" * 48 + record = _make_record(logging.INFO, "response complete") + record.payload = {"api_key": secret, "nested": {"list": [secret]}} + + rendered = JsonFormatter().format(record) + + assert secret not in rendered + assert "REDACTED" in rendered + + +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_diagnostic_filter_redacts_a_non_string_message_object(monkeypatch, native): + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) + secret = "sk-" + "q" * 48 + record = _make_record(logging.ERROR, {"api_key": secret}) + + assert DiagnosticProcessingFilter().filter(record) is True + + assert secret not in record.getMessage() diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index f58ade11d1c..71369c33a6b 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -19,7 +19,11 @@ from litellm._logging import ( verbose_proxy_logger, verbose_router_logger, ) -from litellm.litellm_core_utils.secret_redaction import redact_internal_details, redact_string +from litellm.litellm_core_utils.secret_redaction import ( + redact_internal_details, + redact_string, + redact_structured_value, +) SECRET = "sk-proj-abc123def456ghi789jklmnopqrst" @@ -71,6 +75,17 @@ def test_redact_string_catches_secret_patterns(): assert redact_string(normal) == normal +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +def test_diagnostic_redaction_policy_matches_across_backends(monkeypatch: pytest.MonkeyPatch, native: bool) -> None: + if native: + pytest.importorskip("litellm.rust_bridge._native") + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + + assert redact_string("GET /v1?api_key=abcdefgh12345&page=2") == "GET /v1?REDACTED&page=2" + assert redact_structured_value("db_url", "postgresql://reader@example.org/database") == "REDACTED" + assert redact_internal_details("failed at /etc/service/keys on db.internal") == "failed at REDACTED on REDACTED" + + @pytest.mark.parametrize( "connection_string", [