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:
yucheng-berri 2026-10-17 10:48:30 -07:00
parent e4c0e8b5b5
commit c8cae85463
2 changed files with 61 additions and 8 deletions

View file

@ -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:

View file

@ -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: