fix(guardrails): scan list items beside a capped parent key

List items keep their parent key in the scanned JSON line, capped to its
last 64 characters, so a credential key over a list is still detected
while a long key over many items no longer multiplies the scan text
This commit is contained in:
yucheng-berri 2026-10-17 07:40:59 -07:00
parent 1b6d30fb95
commit 4020206941
4 changed files with 20 additions and 4 deletions

View file

@ -454,6 +454,8 @@ _CONFIG_SECTION: Final = "litellm-prompt"
_JSON_STRING: Final = TypeAdapter(str)
_KEY_CONTEXT: Final = 64
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
@ -546,7 +548,7 @@ def _json_unescaped(value: str) -> str:
def _json_line(key: str | None, text: str) -> str:
value: Final = json.dumps(text, ensure_ascii=False)
return value if key is None else f"{json.dumps(key, ensure_ascii=False)}: {value}"
return value if key is None else f"{json.dumps(key[-_KEY_CONTEXT:], ensure_ascii=False)}: {value}"
def _value_in_texts(value: str, texts: Sequence[str]) -> str | None:

View file

@ -98,12 +98,12 @@ def _tool_arguments(arguments: str) -> JsonValue:
parse_constant=_finite_number,
)
)
json.dumps(parsed, ensure_ascii=False).encode("utf-8")
json.dumps(parsed, ensure_ascii=False).encode("utf-8") # raises for lone surrogates, which stay raw text
return parsed
def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
"""Yield ``(key, text)`` for every non-empty string leaf; the key is None for list items and a bare root."""
"""Yield ``(nearest key, text)`` for every non-empty string leaf; list items take their parent's key."""
pending: Final[list[tuple[object, str | None]]] = [(value, None)] # mutable-ok: owned DFS stack avoids recursion
while pending:
current, key = pending.pop()
@ -120,7 +120,7 @@ def iter_tool_json_fields(value: object) -> Iterator[tuple[str | None, str]]:
items = _object_list(current)
if items is None:
raise ValueError("Invalid tool JSON payload")
pending.extend((child, None) for child in reversed(items))
pending.extend((child, key) for child in reversed(items))
def iter_tool_json_text(value: object) -> Iterator[str]:

View file

@ -214,6 +214,8 @@ def test_read_only_vendors_receive_tool_and_system_text(
def test_moderation_flag_in_tool_result_blocks(rig: Rig) -> None:
rig.wire.drain()
response: Final = rig.proxy.request("POST", "/v1/messages", _messages(FLAGGED))
assert response.status_code == 403, response.text
assert "Violated content safety policy" in response.text
assert _provider_calls(rig.wire.drain()) == ()

View file

@ -147,6 +147,18 @@ async def test_pre_call_hook_records_secrets_found_only_in_tool_payloads():
assert recorded["masked_entity_count"] == {"AWS Access Key": 1}
@pytest.mark.parametrize(
("payload", "expected"),
(
({"password": ["hunter2brahms9x", "safe"]}, {"password": ["[REDACTED]", "safe"]}),
(["a8f5f167f44f4964e6c998dee827110c", "safe"], ["[REDACTED]", "safe"]),
),
ids=("credential-key-over-list", "unkeyed-high-entropy"),
)
def test_list_items_are_scanned_beside_their_parent_key_or_quoted(payload, expected):
assert _guardrail().redact_tool_payload(payload) == expected
def test_long_key_over_a_list_is_scanned_in_bounded_time():
payload = {"k" * 10_000: ["v"] * 1000 + [f"AWS_ACCESS_KEY_ID={AWS_KEY}"]}
started = time.perf_counter()