perf(logging): scan opt-in spend-log tokens natively (#51770)

This commit is contained in:
moe-berri 2026-11-15 06:09:02 +08:00 • committed by GitHub
parent ff997e6970
commit d8e8621b7b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 437 additions and 11 deletions

View file

@ -7,12 +7,15 @@ secrets from strings without depending on the logging-configuration module.
"""
import re
import sys
from functools import lru_cache
from typing import Final
from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH
from litellm.rust_bridge import diagnostics
from litellm.rust_bridge import _native, diagnostics
REDACTED: Final = "REDACTED"
_MINIMUM_TOKEN_LENGTH: Final = min(MINIMUM_CUSTOM_KEY_LENGTH, len("xoxb-") + 10)
_PEM_PRIVATE_KEY_BLOCK: Final = r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----"
@ -138,13 +141,7 @@ _UNMATCHED_JWT: Final = (
_TOKEN_RE: Final = re.compile(f"{_NON_JWT_TOKEN_RE.pattern}|{_TOKEN_START}(?:{_TOKEN_SHAPES[-1]}|{_UNMATCHED_JWT})")
def redact_tokens(value: str, replacement: str = REDACTED) -> str:
"""Replace substrings shaped like a credential (provider key prefixes, JWTs, PEM private keys).
Unlike redact_string() this never matches on surrounding words such as ``password:`` or ``Basic``,
so it is safe to run over free-form prompt and tool content.
"""
def _python_redact_tokens(value: str, replacement: str = REDACTED) -> str:
def replace(match: re.Match[str]) -> str:
if match.lastgroup == "unmatched_jwt":
return _NON_JWT_TOKEN_RE.sub(replacement, match.group())
@ -153,6 +150,26 @@ def redact_tokens(value: str, replacement: str = REDACTED) -> str:
return _TOKEN_RE.sub(replace, value)
@lru_cache(maxsize=1)
def _token_redactor() -> _native.NativeTokenRedactor:
characters: Final = "".join(map(chr, range(128, sys.maxunicode + 1)))
word_ranges: Final = tuple((match.start() + 128, match.end() + 127) for match in re.finditer(r"\w+", characters))
return _native.NativeTokenRedactor(MINIMUM_CUSTOM_KEY_LENGTH, word_ranges)
def redact_tokens(value: str, replacement: str = REDACTED) -> str:
"""Replace credential-shaped tokens without matching surrounding words or field names."""
if "\\" in replacement:
return _python_redact_tokens(value, replacement)
if len(value) < _MINIMUM_TOKEN_LENGTH:
return value
try:
redacted: Final = _token_redactor().redact(value, replacement)
except UnicodeEncodeError:
return _python_redact_tokens(value, replacement)
return value if redacted is None else redacted
_UNIX_SYSTEM_PATH: Final = r"/(?:etc|var|opt|usr|home|root|private|Users|tmp|mnt|srv)/[^\s'\"\)\]}>,]+"
_WINDOWS_DRIVE_PATH: Final = r"[A-Za-z]:\\[^\s'\"\)\]}>,]+"
_PRIVATE_OR_LOOPBACK_IPV4: Final = (

View file

@ -5001,9 +5001,11 @@ dependencies = [
name = "litellm-tracing"
version = "0.1.0"
dependencies = [
"aho-corasick",
"base64 0.22.1",
"fancy-regex 0.19.2",
"percent-encoding",
"regex",
"rstest",
"serde_json",
"tokio",

View file

@ -151,6 +151,8 @@ webpki-roots = "1.0.10"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"
aho-corasick = "1.1.5"
regex = "1.13.1"
veil = "0.3.0"
prost = "0.14.4"
opentelemetry-proto = "0.33"

View file

@ -37,7 +37,7 @@ mod _native {
#[pymodule_export]
use crate::http::unsupported_tls_settings;
#[pymodule_export]
use crate::logger::NativeDiagnosticProcessor;
use crate::logger::{NativeDiagnosticProcessor, NativeTokenRedactor};
#[pymodule_export]
use crate::model_cost::{
NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched, NativeInstalled,
@ -151,6 +151,7 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeTokenRedactor",
"NativeRedisClient",
"NativeRedisBreaker",
"NativeAdmission",

View file

@ -2,14 +2,41 @@ mod machine;
pub(crate) use machine::LoggedMachine;
use litellm_host_python::Pythonized;
use litellm_tracing::{DiagnosticInput, Level, Logger, Metadata, Policy, Processor, Record, Sink};
use litellm_host_python::{Pythonized, release_gil};
use litellm_tracing::{
DiagnosticInput, Level, Logger, Metadata, Policy, Processor, Record, Sink, TokenRedactor,
};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use pyo3::pybacked::PyBackedStr;
const MODULE: &str = "litellm.rust_bridge.host.logger";
type NativeDiagnosticOutput = (String, Option<String>, Option<String>, Vec<String>, bool);
#[pyclass]
pub(crate) struct NativeTokenRedactor {
inner: TokenRedactor,
}
#[pymethods]
impl NativeTokenRedactor {
#[new]
fn new(minimum_custom_key_length: usize, word_ranges: Vec<(u32, u32)>) -> Self {
Self {
inner: TokenRedactor::new(minimum_custom_key_length, word_ranges),
}
}
fn redact(
&self,
py: Python<'_>,
text: PyBackedStr,
replacement: PyBackedStr,
) -> Option<String> {
release_gil(py, || self.inner.redact(&text, &replacement))
}
}
#[pyclass]
pub(crate) struct NativeDiagnosticProcessor {
inner: Processor,

View file

@ -6,9 +6,11 @@ license.workspace = true
repository.workspace = true
[dependencies]
aho-corasick.workspace = true
base64.workspace = true
fancy-regex.workspace = true
percent-encoding.workspace = true
regex.workspace = true
serde_json.workspace = true
tracing.workspace = true
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }

View file

@ -15,10 +15,12 @@ use tracing::{
mod layer;
mod processing;
mod redaction;
mod token_redaction;
pub use layer::sink_layer;
pub use processing::{DiagnosticInput, DiagnosticOutput, Policy, Processor};
pub use redaction::{REDACTED, SecretRedactor};
pub use token_redaction::TokenRedactor;
pub use tracing::{Level, Metadata, debug, error, info, trace, warn};
pub struct ByteChunk<'a>(&'a [u8]);

View file

@ -0,0 +1,292 @@
use aho_corasick::AhoCorasick;
use regex::Regex;
const PREFIXES: &[&str] = &[
"-----BEGIN",
"eyJ",
"ya29.",
"AKIA",
"ASIA",
"sk-",
"sk_live_",
"sk_test_",
"rk_live_",
"rk_test_",
"AIza",
"dapi",
"ghp_",
"gho_",
"ghu_",
"ghs_",
"ghr_",
"github_pat_",
"glpat-",
"lin_api_",
"xoxa-",
"xoxb-",
"xoxp-",
"xoxo-",
"xoxs-",
"xoxr-",
"hf_",
"npm_",
"pypi-AgEIcHlwaS5vcmc",
"SG.",
];
pub struct TokenRedactor {
prefixes: AhoCorasick,
tokens: Regex,
word_ranges: Vec<(u32, u32)>,
}
impl TokenRedactor {
pub fn new(minimum_custom_key_length: usize, word_ranges: Vec<(u32, u32)>) -> Self {
let sk_length = minimum_custom_key_length.saturating_sub(3);
let patterns = [
r"ya29\.[A-Za-z0-9_.~+/-]{20,}",
r"(?:AKIA|ASIA)[0-9A-Z]{16}",
&format!(r"sk-[A-Za-z0-9_-]{{{sk_length},}}"),
r"(?:sk|rk)_(?:live|test)_[A-Za-z0-9]{24,}",
r"AIza[0-9A-Za-z_-]{35}",
r"dapi[0-9a-f]{32}",
r"(?:ghp|gho|ghu|ghs|ghr)_[A-Za-z0-9]{36,}",
r"github_pat_[A-Za-z0-9_]{82}",
r"glpat-[A-Za-z0-9_-]{20,}",
r"lin_api_[A-Za-z0-9]{40}",
r"xox[abposr]-[A-Za-z0-9-]{10,}",
r"hf_[A-Za-z0-9]{34,}",
r"npm_[A-Za-z0-9]{36}",
r"pypi-AgEIcHlwaS5vcmc[A-Za-z0-9_-]{50,}",
r"SG\.[A-Za-z0-9_-]{22}\.[A-Za-z0-9_-]{43}",
];
Self {
prefixes: AhoCorasick::new(PREFIXES).expect("token prefixes compile"),
tokens: Regex::new(&format!(r"\A(?:{})", patterns.join("|")))
.expect("token patterns compile"),
word_ranges,
}
}
pub fn redact(&self, text: &str, replacement: &str) -> Option<String> {
let mut output = None;
let mut copied_until = 0;
let mut failed_jwt_until = 0;
let mut failed_jwt_scope_end = 0;
for found in self.prefixes.find_overlapping_iter(text) {
let start = found.start();
if start < copied_until {
continue;
}
let prefix = PREFIXES[found.pattern().as_usize()];
let end = if prefix == "-----BEGIN" {
pem_end(text, found.end())
} else if !self.starts_token(text, start) {
None
} else if prefix == "eyJ" {
if start < failed_jwt_until {
continue;
}
let header_end = jwt_run_end(text, found.end());
let end = jwt_end(text, found.end(), header_end);
if end.is_none() {
// Every later JWT start in this header shares the same invalid suffix.
failed_jwt_until = header_end;
failed_jwt_scope_end = self.failed_jwt_scope_end(text, start, header_end);
}
end
} else {
let scope = if start < failed_jwt_scope_end {
&text[..failed_jwt_scope_end]
} else {
text
};
self.tokens
.find(&scope[start..])
.map(|token| start + token.end())
.filter(|&end| self.ends_token(scope, end, prefix))
};
if let Some(end) = end {
let result = output.get_or_insert_with(String::new);
result.push_str(&text[copied_until..start]);
result.push_str(replacement);
copied_until = end;
}
}
if let Some(result) = &mut output {
result.push_str(&text[copied_until..]);
}
output
}
fn is_word(&self, character: char) -> bool {
if character.is_ascii() {
return character.is_ascii_alphanumeric() || character == '_';
}
let code = u32::from(character);
let index = self.word_ranges.partition_point(|&(_, end)| end < code);
self.word_ranges
.get(index)
.is_some_and(|&(start, _)| start <= code)
}
fn failed_jwt_scope_end(&self, text: &str, start: usize, header_end: usize) -> usize {
if let Some(marker) = text[start..header_end].find("-----BEGIN") {
return start + marker;
}
let dotted_end = ["ya29", "SG"]
.into_iter()
.find_map(|prefix| {
let position = header_end.checked_sub(prefix.len())?;
if position < start
|| !text[position..].starts_with(prefix)
|| !self.starts_token(text, position)
{
return None;
}
let token = self.tokens.find(&text[position..])?;
let end = position + token.end();
self.ends_token(text, end, if prefix == "SG" { "SG." } else { "ya29." })
.then_some(end)
})
.unwrap_or(header_end);
dotted_end
+ text[dotted_end..]
.chars()
.next()
.filter(|c| !c.is_ascii() || !jwt_byte(*c as u8))
.map_or(0, char::len_utf8)
}
fn starts_token(&self, text: &str, start: usize) -> bool {
let before = &text[..start];
if before.chars().next_back().is_none_or(|c| !self.is_word(c)) {
return true;
}
let bytes = before.as_bytes();
bytes.ends_with(b"\\n")
|| bytes.ends_with(b"\\r")
|| bytes.ends_with(b"\\t")
|| (bytes.len() >= 6
&& bytes[bytes.len() - 6..].starts_with(b"\\u")
&& bytes[bytes.len() - 4..].iter().all(u8::is_ascii_hexdigit))
}
fn ends_token(&self, text: &str, end: usize, prefix: &str) -> bool {
match prefix {
"AKIA" | "ASIA" | "dapi" | "lin_api_" | "npm_" => {
text[end..].chars().next().is_none_or(|c| !self.is_word(c))
}
"SG." => text.as_bytes().get(end).is_none_or(|&b| !jwt_byte(b)),
_ => true,
}
}
}
fn jwt_byte(byte: u8) -> bool {
byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'-'
}
fn jwt_run_end(text: &str, start: usize) -> usize {
start
+ text.as_bytes()[start..]
.iter()
.take_while(|&&byte| jwt_byte(byte))
.count()
}
fn jwt_end(text: &str, header_start: usize, header_end: usize) -> Option<usize> {
let bytes = text.as_bytes();
if header_end - header_start < 10 || bytes.get(header_end) != Some(&b'.') {
return None;
}
let body_start = header_end + 1;
let body_end = jwt_run_end(text, body_start);
if body_start == body_end || bytes.get(body_end) != Some(&b'.') {
return None;
}
let signature_start = body_end + 1;
let signature_end = jwt_run_end(text, signature_start);
(signature_end > signature_start).then_some(signature_end)
}
fn pem_label_end(text: &str, start: usize) -> Option<usize> {
let marker = start + text[start..].find("-----")?;
let label = &text[start..marker];
(label.ends_with("PRIVATE KEY")
&& label
.bytes()
.all(|b| b.is_ascii_uppercase() || b == b' ' || b == b'-'))
.then_some(marker + 5)
}
fn pem_end(text: &str, label_start: usize) -> Option<usize> {
let body_start = pem_label_end(text, label_start)?;
let marker = body_start + text[body_start..].find("-----")?;
text[marker..]
.strip_prefix("-----END")
.and_then(|_| pem_label_end(text, marker + 8))
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::{fixture, rstest};
#[fixture]
fn redactor() -> TokenRedactor {
TokenRedactor::new(
16,
vec![
(u32::from('²'), u32::from('²')),
(u32::from('é'), u32::from('é')),
],
)
}
#[rstest]
#[case::jwt("eyJaaaaaaaaaaaa.x.signature", Some("hidden"))]
#[case::failed_jwt("eyJaaaaaaaaaaaa..signature", None)]
#[case::unicode_word("éghp_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", None)]
#[case::unicode_number("²ghp_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", None)]
#[case::combining_mark(
"\u{301}ghp_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
Some("\u{301}hidden")
)]
#[case::escape(
r"\u2022ghp_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
Some(r"\u2022hidden")
)]
#[case::pem(
"-----BEGIN RSA PRIVATE KEY-----\nbody\n-----END PRIVATE KEY-----",
Some("hidden")
)]
#[case::nested_pem(
"-----BEGIN PRIVATE KEY-----\n-----BEGIN PRIVATE KEY-----\n-----END PRIVATE KEY-----",
Some("-----BEGIN PRIVATE KEY-----\nhidden")
)]
#[case::short_key_before_pem(
"eyJa-sk-short-----BEGIN PRIVATE KEY-----x-----END PRIVATE KEY-----",
Some("eyJa-sk-shorthidden")
)]
fn preserves_token_boundaries(
redactor: TokenRedactor,
#[case] input: &str,
#[case] expected: Option<&str>,
) {
assert_eq!(redactor.redact(input, "hidden").as_deref(), expected);
}
#[rstest]
#[case::missing_body("..signature")]
#[case::missing_signature(".body.")]
#[case::no_dots("")]
fn failed_jwt_runs_do_not_hide_inner_keys(redactor: TokenRedactor, #[case] suffix: &str) {
let run = "eyJaaaaaaaaaa-AKIAAAAAAAAAAAAAAAAA-".repeat(24_000);
let expected = "eyJaaaaaaaaaa-hidden-".repeat(24_000);
assert_eq!(
redactor.redact(&format!("{run}{suffix}"), "hidden"),
Some(format!("{expected}{suffix}"))
);
}
}

View file

@ -569,6 +569,11 @@ class NativeDatabase:
def connect() -> Future[NativeDatabase]: ...
def close(self) -> Future[None]: ...
@final
class NativeTokenRedactor:
def __new__(cls, minimum_custom_key_length: int, word_ranges: Sequence[tuple[int, int]]) -> NativeTokenRedactor: ...
def redact(self, text: str, replacement: str) -> str | None: ...
@final
class NativeDiagnosticProcessor:
def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ...
@ -785,6 +790,7 @@ __all__ = [
"NativeSensitiveRoutePins",
"NativeSpendCounters",
"NativeTableSwapped",
"NativeTokenRedactor",
"NativeUsageRead",
"PostCallSettlement",
"ProcessReservedForForking",

View file

@ -548,6 +548,44 @@ def test_stored_messages_mask_a_github_token_inside_tool_output(gateway: Gateway
assert json.loads(string_value(stored_arguments)) == {"cmd": f"git push https://{REDACTED}@github.com/o/r HEAD"}
def test_malformed_jwt_text_preserves_nested_key_redaction(gateway: Gateway, tmp_path: Path) -> None:
access_key: Final = "AKIA" + "A" * 16
content: Final = (f"eyJaaaaaaaaaa-{access_key}-" * 20) + ".body."
redacted: Final = ((f"eyJaaaaaaaaaa-{REDACTED}-") * 20) + ".body."
response_id: Final = f"chatcmpl-jwt-{uuid4()}"
def respond(request: Request) -> Reply:
assert request.method == "POST"
assert _json_object(request.body)["messages"] == [{"role": "user", "content": content}]
return Reply(body=json.dumps(_chat_completion(response_id, content)).encode())
with (
wire_server(_answering_model_listing(respond)) as wire,
owned_proxy(
gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path, redact_tokens=True), workers=2
) as isolated,
isolated.scenario() as scenario,
):
model: Final = scenario.model(
model="openai/synthetic-redaction-model",
api_base=wire.url,
api_key="synthetic-openai-key",
input_cost_per_token=0.001,
output_cost_per_token=0.002,
)
response: Final = isolated.request(
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": content}]}
)
assert response.status_code == 200, response.text
caller_choice: Final = _objects(_json_object(response.content)["choices"])[0]
assert object_value(caller_choice["message"])["content"] == content
assert len(_provider_calls(wire.drain())) == 1
row: Final = _stored_row(response_id)
assert object_value(row["proxy_server_request"])["messages"] == [{"role": "user", "content": redacted}]
stored_choice: Final = _objects(object_value(row["response"])["choices"])[0]
assert object_value(stored_choice["message"])["content"] == redacted
def test_previous_response_id_replay_sends_real_tool_payloads(gateway: Gateway, tmp_path: Path) -> None:
function_arguments: Final = {"sort_key": "created_at", "access_level": "admin"}
function_output: Final = {

View file

@ -94,12 +94,26 @@ def test_token_redaction_masks_known_token_shapes(token: str) -> None:
"-----BEGIN PRIVATE KEY-----\nREDACTED",
),
("-----BEGIN BAD-----\n-----END PRIVATE KEY-----", "-----BEGIN BAD-----\n-----END PRIVATE KEY-----"),
(
"eyJa-sk-short-----BEGIN PRIVATE KEY-----x-----END PRIVATE KEY-----",
"eyJa-sk-shortREDACTED",
),
),
)
def test_token_redaction_preserves_matching_rules(text: str, expected: str) -> None:
assert redact_tokens(text) == expected
@pytest.mark.parametrize("character", ("é", "²", "\u0301", "\U00014153", "\u1c89"))
def test_token_boundaries_follow_python_unicode_version(character: str) -> None:
token: Final = "ghp_" + "A" * 36
key: Final = "AKIA" + "A" * 16
expected_token: Final = token if character.isalnum() else "REDACTED"
expected_key: Final = key if character.isalnum() else "REDACTED"
assert redact_tokens(character + token) == character + expected_token
assert redact_tokens(key + character) == expected_key + character
@pytest.mark.parametrize(
"content,expected",
(
@ -122,6 +136,29 @@ def test_failed_jwt_header_does_not_hide_other_tokens(content: str, expected: st
assert redact_tokens(prefix + content) == prefix + expected
@pytest.mark.parametrize("suffix", ("", "..signature", ".body."))
def test_failed_jwt_run_masks_each_embedded_bounded_key(suffix: str) -> None:
run: Final = "eyJaaaaaaaaaa-AKIA" + "A" * 16 + "-"
expected: Final = "eyJaaaaaaaaaa-REDACTED-"
assert redact_tokens(run * 24_000 + suffix) == expected * 24_000 + suffix
@pytest.mark.parametrize(
"text, replacement, expected",
(
("\u0301sk-" + "a" * 20, "hidden", "\u0301hidden"),
("²sk-" + "a" * 20, "hidden", "²sk-" + "a" * 20),
("\ud800sk-" + "a" * 20, "hidden", "\ud800hidden"),
("sk-" + "a" * 20, r"\n", "\n"),
("sk-" + "a" * 20, "", ""),
),
)
def test_token_redaction_preserves_unicode_and_replacement_semantics(
text: str, replacement: str, expected: str
) -> None:
assert redact_tokens(text, replacement) == expected
@pytest.fixture(autouse=True)
def _enable_redaction():
"""Ensure secret redaction is on (the default) for all tests in this module."""