diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 88c48ea0b8d..40f25f6b36d 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -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[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P\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: diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index b095c853ebe..c352545ab81 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -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]: diff --git a/tests/integration/security/test_tool_text_guardrail_consumers.py b/tests/integration/security/test_tool_text_guardrail_consumers.py index b4d7ff79385..4a5c8afbbb1 100644 --- a/tests/integration/security/test_tool_text_guardrail_consumers.py +++ b/tests/integration/security/test_tool_text_guardrail_consumers.py @@ -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()) == () diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index 7cb38f214db..f398046f602 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py @@ -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()