mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
perf(logging): scan opt-in spend-log tokens natively (#51770)
This commit is contained in:
parent
ff997e6970
commit
d8e8621b7b
11 changed files with 437 additions and 11 deletions
|
|
@ -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 = (
|
||||
|
|
|
|||
2
litellm/rust/Cargo.lock
generated
2
litellm/rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
|
|
|
|||
292
litellm/rust/crates/tracing/src/token_redaction.rs
Normal file
292
litellm/rust/crates/tracing/src/token_redaction.rs
Normal 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}"))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue