mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(guardrails): size the clean-scan cache for many sessions and replace longest secrets first
4096 digests held about 20 long agent sessions per worker, so interleaved sessions evicted each other and every resend rescanned its history; 65536 digests (about 10 MB) hold a few hundred. Tool payload secrets are now replaced longest first after JSON unescaping, so a shorter overlapping match cannot leave part of a secret behind. The key-context test checks the work it bounds instead of wall-clock time, which CI could not meet
This commit is contained in:
parent
e4c0e8b5b5
commit
c8cae85463
2 changed files with 61 additions and 8 deletions
|
|
@ -468,8 +468,9 @@ _SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
|
|||
|
||||
_SCAN_SUFFIX: Final = ".py"
|
||||
|
||||
# Digests of scan inputs that held no secret, so the history an agent resends every turn is scanned once.
|
||||
_CLEAN_SCAN_CACHE_SIZE: Final = 4096
|
||||
# Digests of scan inputs that held no secret, so the history an agent resends every turn is scanned once. 65536
|
||||
# digests take about 10 MB and hold a few hundred interleaved agent sessions per worker.
|
||||
_CLEAN_SCAN_CACHE_SIZE: Final = 65536
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
|
|
@ -644,7 +645,8 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
|
|||
"Detected and redacted secrets in tool payload: %s",
|
||||
sorted(dict.fromkeys(secret["type"] for secret in detected)),
|
||||
)
|
||||
pattern: Final = re.compile("|".join(re.escape(secret["value"]) for secret in detected))
|
||||
values: Final = sorted({secret["value"] for secret in detected}, key=len, reverse=True) # longest match wins
|
||||
pattern: Final = re.compile("|".join(re.escape(value) for value in values))
|
||||
return map_tool_json_text(payload, lambda text: pattern.sub("[REDACTED]", text))
|
||||
|
||||
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
|
|
|
|||
|
|
@ -160,6 +160,40 @@ def test_clean_scan_reuse_is_bounded(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
assert scanned == ["first", "second", "third", "second", "first"]
|
||||
|
||||
|
||||
def test_interleaved_agent_sessions_reuse_their_clean_scans(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
scanned: Final[list[str]] = []
|
||||
|
||||
def clean_scan(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
|
||||
scanned.append(lines[0])
|
||||
return frozenset()
|
||||
|
||||
monkeypatch.setattr(secret_detection, "_scan_lines", clean_scan)
|
||||
guardrail: Final = _guardrail()
|
||||
sessions: Final = [[f"session {session} turn {turn}" for turn in range(201)] for session in range(25)]
|
||||
for history in sessions:
|
||||
for text in history:
|
||||
guardrail.redact_text(text)
|
||||
scanned.clear()
|
||||
for history in sessions:
|
||||
for text in history:
|
||||
guardrail.redact_text(text)
|
||||
assert scanned == []
|
||||
|
||||
|
||||
def test_overlapping_secrets_are_replaced_longest_first(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
guardrail: Final = _guardrail()
|
||||
detected: Final = [
|
||||
{"type": "Secret Keyword", "value": "\\u0041\\u0042"},
|
||||
{"type": "Secret Keyword", "value": "ABCD"},
|
||||
]
|
||||
|
||||
def overlapping_scan(text: str) -> list[dict[str, str]]:
|
||||
return detected
|
||||
|
||||
monkeypatch.setattr(guardrail, "scan_message_for_secrets", overlapping_scan)
|
||||
assert guardrail.redact_tool_payload({"token": "ABCDxyz"}) == {"token": "[REDACTED]xyz"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_records_secrets_found_only_in_tool_payloads():
|
||||
data = {
|
||||
|
|
@ -202,12 +236,29 @@ def test_list_items_are_scanned_beside_their_parent_key_or_quoted(payload, expec
|
|||
(("k" * 10_000, 1000), ("\u0000" * 64, 3000), ("tok " * 250_000, 2000)),
|
||||
ids=("long", "control-characters", "long-over-many-items"),
|
||||
)
|
||||
def test_keys_over_a_list_are_scanned_in_bounded_time(key, items):
|
||||
payload = {key: ["v"] * items + [f"AWS_ACCESS_KEY_ID={AWS_KEY}"]}
|
||||
started = time.perf_counter()
|
||||
result = _guardrail().redact_tool_payload(payload)
|
||||
assert time.perf_counter() - started < 3.0
|
||||
def test_keys_over_a_list_build_one_bounded_context(monkeypatch: pytest.MonkeyPatch, key: str, items: int) -> None:
|
||||
contexts: Final[list[str]] = []
|
||||
scans: Final[list[str]] = []
|
||||
key_context: Final = secret_detection._key_context
|
||||
guardrail: Final = _guardrail()
|
||||
scan: Final = guardrail.scan_message_for_secrets
|
||||
|
||||
def recorded_context(value: str) -> str:
|
||||
contexts.append(value)
|
||||
return key_context(value)
|
||||
|
||||
def recorded_scan(text: str) -> list[dict[str, str]]:
|
||||
scans.append(text)
|
||||
return scan(text)
|
||||
|
||||
monkeypatch.setattr(secret_detection, "_key_context", recorded_context)
|
||||
monkeypatch.setattr(guardrail, "scan_message_for_secrets", recorded_scan)
|
||||
secret_line: Final = f"AWS_ACCESS_KEY_ID={AWS_KEY}"
|
||||
result: Final = guardrail.redact_tool_payload({key: ["v"] * items + [secret_line]})
|
||||
assert result == {key: ["v"] * items + ["AWS_ACCESS_KEY_ID=[REDACTED]"]}
|
||||
assert contexts == [key]
|
||||
longest: Final = secret_detection._KEY_CONTEXT + len('"": ') + len(json.dumps(secret_line))
|
||||
assert max(len(line) for line in scans[0].splitlines()) == longest
|
||||
|
||||
|
||||
def _recorded(request_data: dict) -> dict:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue