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.

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, ...]:

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 = {