From acb1b9e49ccc1a9c8429bf1b77820115d549b835 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 19:19:52 -0700 Subject: [PATCH] fix(guardrails): keep both scans' masks on an escaped systemone field A state entry with a JSON escape is scanned twice, once as the escaped entry with its key as context and once decoded. Taking only the decoded rewrite dropped a mask that needed the key, so a tax number under a TFN key reached the provider when an email in the same value was also masked. Both rewrites' edits now apply together, word by word, and a larger mask wins over a smaller one inside it. Edits that partly overlap reject the request as unappliable. The entry scan sees each letter escape followed by a tilde that is removed afterwards, so a match next to an escape no longer swallows the escape letter. --- .../guardrail_translation/handler.py | 73 +++++++++++++--- .../test_decisions_guardrail_handler.py | 83 +++++++++++++++++-- 2 files changed, 139 insertions(+), 17 deletions(-) diff --git a/litellm/llms/openai/decisions/guardrail_translation/handler.py b/litellm/llms/openai/decisions/guardrail_translation/handler.py index ead6f3e176c..5b02a037041 100644 --- a/litellm/llms/openai/decisions/guardrail_translation/handler.py +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -1,9 +1,10 @@ import json import re from collections.abc import Iterator, Mapping +from difflib import SequenceMatcher from itertools import chain, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, ClassVar, Final, TypeVar +from typing import TYPE_CHECKING, ClassVar, Final from pydantic import JsonValue, TypeAdapter, ValidationError @@ -21,8 +22,10 @@ if TYPE_CHECKING: _JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _KEY_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) _JSON_KEY_PREFIX: Final = re.compile(r'("(?:[^"\\\x00-\x1f]|\\(?:["\\/bfnrt]|u[0-9a-fA-F]{4}))*"): ') +_JSON_ESCAPE: Final = re.compile(r"(\\(?:[bfnrt]|u[0-9a-fA-F]{4}))|\\.") +_MARKED_JSON_ESCAPE: Final = re.compile(r"(\\(?:[bfnrt]|u[0-9a-fA-F]{4}))~|\\.") +_TOKEN: Final = re.compile(r"\s+|\S+") _DESCRIBED_ITEM_KEYS: Final = ("choices", "levels") -_Field: Final = TypeVar("_Field", bound=JsonValue) _QUESTION_LIST_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, JsonValue], ...]]] = TypeAdapter( tuple[Mapping[str, JsonValue], ...] ) @@ -81,8 +84,16 @@ def _guardrailed_input(decision_input: JsonValue, guardrailed: Mapping[str, str] return [_guardrailed_message(message, guardrailed) for message in decision_input] +def _marked_escape(escape: re.Match[str]) -> str: + return f"{escape.group()}~" if escape.group(1) else escape.group() + + +def _unmarked_escape(escape: re.Match[str]) -> str: + return escape.group(1) or escape.group() + + def _entry_text(key: str, value: str | float | None) -> str: - return json.dumps({key: value}, ensure_ascii=False)[1:-1] + return _JSON_ESCAPE.sub(_marked_escape, json.dumps({key: value}, ensure_ascii=False)[1:-1]) def _escaped_texts(text: str, key: str, item: JsonValue) -> tuple[str, ...]: @@ -164,18 +175,57 @@ def _rewritten_value(text: str) -> JsonValue: def _rewritten_entry(rewritten: str, guardrail_name: str | None) -> tuple[str, JsonValue]: - key: Final = _JSON_KEY_PREFIX.match(rewritten) + unmarked: Final = _MARKED_JSON_ESCAPE.sub(_unmarked_escape, rewritten) + key: Final = _JSON_KEY_PREFIX.match(unmarked) if key is None: raise unappliable_request_rewrite(guardrail_name) - return _KEY_ADAPTER.validate_json(key.group(1)), _rewritten_value(rewritten[key.end() :]) + return _KEY_ADAPTER.validate_json(key.group(1)), _rewritten_value(unmarked[key.end() :]) def _decoded_text(text: str, form: str, guardrailed: Mapping[str, str]) -> str: return form if form in text else guardrailed.get(form, form) -def _merged_field(original: _Field, rewritten: _Field, decoded: _Field) -> _Field: - return rewritten if decoded == original else decoded +def _tokens(text: str) -> tuple[str, ...]: + return tuple(_TOKEN.findall(text)) + + +def _edits(original: tuple[str, ...], rewritten: tuple[str, ...]) -> frozenset[tuple[int, int, str]]: + opcodes: Final = SequenceMatcher(None, original, rewritten, autojunk=False).get_opcodes() + return frozenset( + (start, end, "".join(rewritten[low:high])) for tag, start, end, low, high in opcodes if tag != "equal" + ) + + +def _covered(edit: tuple[int, int, str], edits: frozenset[tuple[int, int, str]]) -> bool: + start, end, _ = edit + return any(other[0] <= start and end <= other[1] and other[1] - other[0] > end - start for other in edits) + + +def _merged_text(original: str, rewritten: str, decoded: str, guardrail_name: str | None) -> str: + if decoded in (original, rewritten): + return rewritten + if rewritten == original: + return decoded + tokens: Final = _tokens(original) + edits: Final = _edits(tokens, _tokens(rewritten)) | _edits(tokens, _tokens(decoded)) + kept: Final = tuple(sorted(edit for edit in edits if not _covered(edit, edits))) + if any(later[0] < earlier[1] for earlier, later in zip(kept, kept[1:])): + raise unappliable_request_rewrite(guardrail_name) + kept_starts: Final = (0, *(end for _, end, _ in kept)) + kept_ends: Final = (*(start for start, _, _ in kept), len(tokens)) + replacements: Final = (*(text for _, _, text in kept), "") + return "".join( + "".join(tokens[start:end]) + text for start, end, text in zip(kept_starts, kept_ends, replacements, strict=True) + ) + + +def _merged_value(original: str, rewritten: JsonValue, decoded: str, guardrail_name: str | None) -> JsonValue: + if decoded == original: + return rewritten + if isinstance(rewritten, str): + return _merged_text(original, rewritten, decoded, guardrail_name) + raise unappliable_request_rewrite(guardrail_name) def _guardrailed_entry( @@ -186,11 +236,10 @@ def _guardrailed_entry( text: Final = _entry_text(key, item) rewritten: Final = guardrailed.get(text, text) entry: Final = (key, item) if rewritten == text else _rewritten_entry(rewritten, guardrail_name) - decoded_item: Final = _decoded_text(text, item, guardrailed) if isinstance(item, str) else item - return ( - _merged_field(key, entry[0], _decoded_text(text, key, guardrailed)), - _merged_field(item, entry[1], decoded_item), - ) + merged_key: Final = _merged_text(key, entry[0], _decoded_text(text, key, guardrailed), guardrail_name) + if not isinstance(item, str): + return merged_key, entry[1] + return merged_key, _merged_value(item, entry[1], _decoded_text(text, item, guardrailed), guardrail_name) def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: diff --git a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py index a7e5fa3711f..d971f750d4c 100644 --- a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py +++ b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py @@ -1,7 +1,9 @@ +import re from functools import reduce from typing import Final, Literal import pytest +from litellm_enterprise.enterprise_callbacks.secret_detection import _ENTERPRISE_SecretDetection from pydantic import JsonValue from litellm.integrations.custom_guardrail import CustomGuardrail @@ -63,6 +65,32 @@ class _DropsQuotes(CustomGuardrail): return GenericGuardrailAPIInputs(texts=[text.replace('"', "") for text in inputs.get("texts", [])]) +class _MasksOverlappingPairs(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=[text.replace("b\nc", "[BC]").replace("a b", "[AB]") for text in inputs.get("texts", [])] + ) + + +class _NullsAMaskedValue(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs( + texts=[re.sub(r'": ".*"$', '": null', text).replace("b", "[B]") for text in inputs.get("texts", [])] + ) + + class _ReturnsNoTexts(CustomGuardrail): async def apply_guardrail( self, @@ -234,29 +262,42 @@ async def test_systemone_masks_a_card_number_that_json_escaping_splits() -> None @pytest.mark.asyncio @pytest.mark.parametrize( - ("pattern_name", "state", "masked"), + ("pattern_names", "state", "masked"), ( pytest.param( - "visa", + ("visa",), {"4111111111111111": "4111\t1111\t1111\t1111"}, {"[VISA_REDACTED]": "[VISA_REDACTED]"}, id="key-and-escaped-value", ), pytest.param( - "email", + ("email",), {"note": "Please call me back.\ndana.quill@example.com"}, {"note": "Please call me back.\n[EMAIL_REDACTED]"}, id="match-after-an-escape", ), + pytest.param( + ("au_tfn", "email"), + {"TFN": "123456789 contact dana@example.com\n"}, + {"TFN": "[AU_TFN_REDACTED] contact [EMAIL_REDACTED]\n"}, + id="key-context-match-in-an-escaped-value", + ), + pytest.param( + ("au_tfn", "visa"), + {"TFN": "123456789\n4111\t1111\t1111\t1111"}, + {"TFN": "[AU_TFN_REDACTED]\n[VISA_REDACTED]"}, + id="key-context-match-beside-a-split-match", + ), ), ) async def test_systemone_masks_findings_on_both_sides_of_a_json_escape( - pattern_name: str, state: dict[str, JsonValue], masked: dict[str, JsonValue] + pattern_names: tuple[str, ...], state: dict[str, JsonValue], masked: dict[str, JsonValue] ) -> None: guardrail: Final = ContentFilterGuardrail( guardrail_name="mask", patterns=[ - ContentFilterPattern(pattern_type="prebuilt", pattern_name=pattern_name, action=ContentFilterAction.MASK) + ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK) + for name in pattern_names ], ) data: Final = {"model": "judge", "state": state, "questions": {"q": {"type": "noul", "criteria": {}}}} @@ -266,6 +307,38 @@ async def test_systemone_masks_findings_on_both_sides_of_a_json_escape( assert result["state"] == masked +@pytest.mark.asyncio +async def test_systemone_keeps_a_secret_mask_that_covers_an_escaped_value() -> None: + data: Final = { + "model": "judge", + "state": {"password": "hunter2\nAKIAIOSFODNN7EXAMPLE"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages( + data, _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets") + ) + + assert result["state"] == {"password": "[REDACTED]"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guardrail", "note"), + ( + pytest.param(_MasksOverlappingPairs(guardrail_name="pairs"), "a b\nc", id="partly-overlapping-masks"), + pytest.param(_NullsAMaskedValue(guardrail_name="nulls"), "a\nb", id="value-turned-into-null"), + ), +) +async def test_systemone_rejects_an_escaped_entry_whose_two_rewrites_cannot_merge( + guardrail: CustomGuardrail, note: str +) -> None: + data: Final = {"model": "judge", "state": {"note": note}, "questions": {"q": {"type": "noul", "criteria": {}}}} + + with pytest.raises(UnappliableRequestRewrite): + await SystemOneHandler().process_input_messages(data, guardrail) + + @pytest.mark.asyncio async def test_systemone_keeps_every_state_entry_when_masked_keys_collide() -> None: data: Final = {