Merge pull request #41952 from BerriAI/litellm_masker_memo_depth_fail_closed

fix(masker): memoize shared nodes and fail closed past the depth cap

litellm/litellm_core_utils/sensitive_data_masker.py

@@ -1,5 +1,6 @@
 from collections.abc import Mapping, Sequence
 from collections.abc import Set as AbstractSet
+from dataclasses import dataclass, field
 from typing import Any, Final
 
 from pydantic import BaseModel
@@ -176,26 +177,49 @@ def mask_credentials_in_payload(data: object) -> object:
     config-dump semantics (``None`` -> ``"None"``, tuples stringified,
     objects flattened via ``__dict__``) would silently distort the record.
 
+    A container referenced from several places in ``data`` is rebuilt once and
+    referenced from the same places in the copy, so a shared subtree never
+    fans out into independent copies, and a reference back into a container
+    still being rebuilt (a cycle) becomes ``REDACTED``. A container nested past
+    ``DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER`` is replaced by
+    ``REDACTED`` rather than returned unmasked.
+
     Sensitive-key detection is delegated to the shared
     :class:`SensitiveDataMasker` so pattern updates stay in one place.
     """
-    return _walk_payload(data, key_is_sensitive=False, depth=0)
+    return _PayloadWalker().walk(data, key_is_sensitive=False, depth=0)
 
 
-def _walk_payload(node: object, key_is_sensitive: bool, depth: int) -> object:
-    if depth >= DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER:
-        return node
-    if isinstance(node, Mapping):
-        return {k: _walk_payload(v, _default_masker.is_sensitive_key(k), depth + 1) for k, v in node.items()}
-    if isinstance(node, list):
-        return [_walk_payload(item, key_is_sensitive, depth + 1) for item in node]
-    if isinstance(node, tuple):
-        return tuple(_walk_payload(item, key_is_sensitive, depth + 1) for item in node)
-    if isinstance(node, BaseModel):
-        return _walk_payload(node.model_dump(), key_is_sensitive, depth)
-    if key_is_sensitive and isinstance(node, str) and node:
-        return _default_masker._mask_value(node)
-    return node
+@dataclass(frozen=True, slots=True)
+class _PayloadWalker:
+    _memo: dict[tuple[int, bool], tuple[object, object]] = field(  # mutable-ok: memo of one walk, pins each keyed node
+        default_factory=dict
+    )
+
+    def walk(self, node: object, key_is_sensitive: bool, depth: int) -> object:
+        if not isinstance(node, (Mapping, list, tuple, BaseModel)):
+            return _default_masker._mask_value(node) if key_is_sensitive and isinstance(node, str) and node else node
+        if depth >= DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER:
+            return REDACTED
+        memo_key: Final = (id(node), key_is_sensitive and not isinstance(node, Mapping))
+        cached: Final = self._memo.get(memo_key)
+        if cached is not None:
+            return cached[1]
+        self._memo[memo_key] = (node, REDACTED)
+        rebuilt: Final = self._rebuild(node, key_is_sensitive, depth)
+        self._memo[memo_key] = (node, rebuilt)
+        return rebuilt
+
+    def _rebuild(
+        self, node: Mapping[str, object] | Sequence[object] | BaseModel, key_is_sensitive: bool, depth: int
+    ) -> object:
+        if isinstance(node, BaseModel):
+            return self.walk(node.model_dump(), key_is_sensitive, depth)
+        if isinstance(node, Mapping):
+            return {k: self.walk(v, _default_masker.is_sensitive_key(k), depth + 1) for k, v in node.items()}
+        if isinstance(node, tuple):
+            return tuple(self.walk(item, key_is_sensitive, depth + 1) for item in node)
+        return [self.walk(item, key_is_sensitive, depth + 1) for item in node]
 
 
 def mask_sensitive_keys(data: dict[str, Any], sensitive_fields: set[str]) -> dict[str, Any]:

tests/code_coverage_tests/recursive_detector.py

@@ -36,7 +36,6 @@ IGNORE_FUNCTIONS = [
     "_collect_argument_paths",  # max depth set.
     "_split_text",  # max depth set.
     "_mask_sequence",  # max depth set.
-    "_walk_payload",  # max depth set (DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER).
     "_delete_nested_value_custom",  # max depth set (bounded by number of path segments).
     "filter_exceptions_from_params",  # max depth set (default 20) to prevent infinite recursion.
     "__getattr__",  # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion.

tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py

@@ -2,11 +2,11 @@
 Unit tests for SensitiveDataMasker - List Preservation
 """
 
+from functools import reduce
+from typing import Final
 
 import pytest
 
-# Add the parent directory to the system path
-
 from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
 
 
@@ -152,9 +152,7 @@ def test_mask_short_values_false_keeps_short_values_readable():
     chars of an exception and only masks longer tails), while longer values are still
     partially masked.
     """
-    masker = SensitiveDataMasker(
-        visible_prefix=50, visible_suffix=0, mask_short_values=False
-    )
+    masker = SensitiveDataMasker(visible_prefix=50, visible_suffix=0, mask_short_values=False)
 
     short = "Test exception for structure validation"
     assert masker._mask_value(short) == short
@@ -202,9 +200,7 @@ def test_mask_sensitive_structure_passes_through_plain_topology_names():
     from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
 
     assert mask_sensitive_structure(["gpt-4", "claude-3-haiku"]) == ["gpt-4", "claude-3-haiku"]
-    assert mask_sensitive_structure([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) == [
-        {"gpt-3.5-turbo": ["claude-3-haiku"]}
-    ]
+    assert mask_sensitive_structure([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) == [{"gpt-3.5-turbo": ["claude-3-haiku"]}]
     assert mask_sensitive_structure(None) is None
 
 
@@ -233,9 +229,7 @@ def test_mask_sensitive_structure_masks_credentials_nested_in_config_shape():
     from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
 
     secret = "sk-NESTEDINLINESECRET0987654321"
-    masked = mask_sensitive_structure(
-        [{"primary-group": [{"model": "gpt-4o", "api_key": secret}]}]
-    )
+    masked = mask_sensitive_structure([{"primary-group": [{"model": "gpt-4o", "api_key": secret}]}])
     assert secret not in str(masked)
 
 
@@ -282,10 +276,7 @@ def test_mask_credentials_in_payload_masks_inside_pydantic_models():
     auth_dict = result["user_api_key_auth"]
     assert isinstance(auth_dict, dict)
     assert auth_dict["team_alias"] == "acme"
-    assert (
-        auth_dict["token"]
-        != "1b01552f6e52e0d41963dd6a185bd6b074624e330999534ca7ff5adfdf622dfc"
-    )
+    assert auth_dict["token"] != "1b01552f6e52e0d41963dd6a185bd6b074624e330999534ca7ff5adfdf622dfc"
     assert "*" in auth_dict["token"]
 
 
@@ -314,6 +305,159 @@ def test_mask_credentials_in_payload_masks_only_sensitive_string_leaves():
     assert masked.endswith(plaintext[-4:])
 
 
+def _unique_dict_ids(node: object) -> frozenset[int]:
+    if isinstance(node, dict):
+        return frozenset((id(node),)).union(*(_unique_dict_ids(value) for value in node.values()))
+    if isinstance(node, list):
+        return frozenset().union(*(_unique_dict_ids(value) for value in node))
+    return frozenset()
+
+
+def _nested_under_levels(leaf: object, levels: int) -> object:
+    return reduce(lambda inner, level: {f"l{level}": inner}, range(levels, 0, -1), leaf)
+
+
+def test_mask_credentials_in_payload_keeps_a_shared_dict_shared():
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    shared: Final = {"api_key": "sk-shared-1234567890abcdef", "model": "gpt-4o-mini"}
+    result: Final = mask_credentials_in_payload({"first": shared, "second": shared})
+
+    assert result["first"] is result["second"]
+    assert result["first"]["model"] == "gpt-4o-mini"
+    assert result["first"]["api_key"] != "sk-shared-1234567890abcdef"
+
+
+def test_mask_credentials_in_payload_walks_each_dag_node_once():
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    root: Final = reduce(
+        lambda inner, _: {"a": inner, "b": inner, "c": inner}, range(8), {"api_key": "sk-leaf-1234567890abcdef"}
+    )
+
+    result: Final = mask_credentials_in_payload(root)
+
+    assert len(_unique_dict_ids(root)) == 9
+    assert len(_unique_dict_ids(result)) == 9
+    assert "sk-leaf-1234567890abcdef" not in str(result)
+
+
+def test_mask_credentials_in_payload_cuts_a_cycle_at_its_first_back_edge():
+    from litellm.litellm_core_utils.secret_redaction import REDACTED
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    node: Final[dict[str, object]] = {"api_key": "sk-cycle-1234567890abcdef"}
+    node["kids"] = [node] * 3
+
+    result: Final = mask_credentials_in_payload(node)
+
+    assert result["kids"] == [REDACTED, REDACTED, REDACTED]
+    assert result["api_key"] != "sk-cycle-1234567890abcdef"
+    assert len(_unique_dict_ids(result)) == 1
+
+
+def test_mask_credentials_in_payload_masks_a_shared_list_only_under_a_sensitive_key():
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    shared: Final = ["sk-list-1234567890abcdef"]
+
+    plain_first: Final = mask_credentials_in_payload({"tags": shared, "api_key": shared})
+    assert plain_first["tags"] == ["sk-list-1234567890abcdef"]
+    assert plain_first["api_key"] != ["sk-list-1234567890abcdef"]
+
+    sensitive_first: Final = mask_credentials_in_payload({"api_key": shared, "tags": shared})
+    assert sensitive_first["api_key"] != ["sk-list-1234567890abcdef"]
+    assert sensitive_first["tags"] == ["sk-list-1234567890abcdef"]
+
+
+def test_mask_credentials_in_payload_masks_a_shared_root_model_list_only_under_a_sensitive_key():
+    from pydantic import RootModel
+
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    shared: Final = RootModel[list[str]](["sk-root-1234567890abcdef"])
+
+    plain_first: Final = mask_credentials_in_payload({"tags": shared, "api_key": shared})
+    assert plain_first["tags"] == ["sk-root-1234567890abcdef"]
+    assert plain_first["api_key"] != ["sk-root-1234567890abcdef"]
+
+    sensitive_first: Final = mask_credentials_in_payload({"api_key": shared, "tags": shared})
+    assert sensitive_first["api_key"] != ["sk-root-1234567890abcdef"]
+    assert sensitive_first["tags"] == ["sk-root-1234567890abcdef"]
+
+
+def test_mask_credentials_in_payload_masks_a_root_model_string_as_one_string():
+    from pydantic import RootModel
+
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    result: Final = mask_credentials_in_payload(
+        {"api_key": RootModel[str]("sk-root-1234567890abcdef"), "model": RootModel[str]("gpt-5.4-mini")}
+    )
+
+    assert result["model"] == "gpt-5.4-mini"
+    assert result["api_key"] != "sk-root-1234567890abcdef"
+    assert result["api_key"].startswith("sk-r")
+
+
+def test_mask_credentials_in_payload_hides_containers_past_the_depth_cap():
+    from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
+    from litellm.litellm_core_utils.secret_redaction import REDACTED
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    secret: Final = "sk-deep-1234567890abcdef"
+    cap: Final = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
+
+    result: Final = mask_credentials_in_payload(_nested_under_levels({"api_key": secret}, cap))
+
+    assert secret not in str(result)
+    at_cap: Final = reduce(lambda node, level: node[f"l{level}"], range(1, cap), result)
+    assert at_cap == {f"l{cap}": REDACTED}
+
+
+def test_mask_credentials_in_payload_treats_strings_at_the_depth_cap_per_key():
+    from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    secret: Final = "sk-deep-1234567890abcdef"
+    cap: Final = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
+
+    strings_at_cap: Final = reduce(
+        lambda node, level: node[f"l{level}"],
+        range(1, cap),
+        mask_credentials_in_payload(_nested_under_levels({"api_key": secret, "model": "gpt-5.4-mini"}, cap - 1)),
+    )
+    assert strings_at_cap["model"] == "gpt-5.4-mini"
+    assert strings_at_cap["api_key"] != secret
+    assert strings_at_cap["api_key"].startswith("sk-d")
+
+
+def test_mask_credentials_in_payload_keeps_sibling_models_apart():
+    """CPython reuses a freed temporary's id, so an id-keyed memo has to pin what it keys."""
+    from pydantic import BaseModel
+
+    from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
+
+    class Inner(BaseModel):
+        label: str
+        api_key: str
+
+    class Outer(BaseModel):
+        inner: Inner
+
+    result: Final = mask_credentials_in_payload(
+        {
+            "first": Outer(inner=Inner(label="one", api_key="sk-first-1234567890abcdef")),
+            "second": Outer(inner=Inner(label="two", api_key="sk-second-1234567890abcdef")),
+        }
+    )
+
+    assert result["first"]["inner"]["label"] == "one"
+    assert result["second"]["inner"]["label"] == "two"
+    assert "sk-second-1234567890abcdef" not in str(result)
+    assert result["second"]["inner"]["api_key"].startswith("sk-s")
+
+
 def test_extra_sensitive_patterns_add_to_the_defaults():
     from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
 
@@ -344,6 +488,7 @@ def test_the_second_positional_argument_is_still_the_override_set():
     assert masker.is_sensitive_key("session_token") is False
     assert masker.is_sensitive_key("auth_token") is True
 
+
 def test_redact_credentials_in_payload_leaves_no_fragment_of_the_secret():
     """A payload rendered straight to stdout cannot afford the partial reveal
     mask_credentials_in_payload leaves, so every credential-named value is replaced