From 7edafd17150a2732b662b69353f90dbdc9419e99 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:28:50 -0700 Subject: [PATCH 1/3] fix(masker): memoize shared nodes and fail closed past the depth cap --- .../sensitive_data_masker.py | 52 ++++-- .../test_sensitive_data_masker.py | 156 ++++++++++++++++-- 2 files changed, 178 insertions(+), 30 deletions(-) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index b4c1beea33e..b68c97e18c9 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/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,47 @@ 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. 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] + 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._rebuild(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]: diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index 1551c3fd6e6..de511b0ce11 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/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,140 @@ 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(): + """One dict referenced twice comes back as one masked dict referenced + twice. Rebuilding each reference separately is what turned an aliased + retry breadcrumb graph exponential in the v1.100.0 OOM.""" + 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(): + """A DAG of 9 dicts where every level references the level below three + times stays 9 dicts after masking, instead of fanning out to 3**8.""" + 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_masks_a_shared_list_only_under_a_sensitive_key(): + """The same list reached under a plain key and under a sensitive key is + masked in the sensitive spot only, whichever reference the walk meets + first, so the memo can neither leak a secret nor mask a plain value.""" + 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(): + """A pydantic model that dumps to a list is a list once walked, so the + memo must keep its plain and sensitive rebuilds apart the same way, or + the reference met first decides what the other one shows.""" + 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_hides_containers_past_the_depth_cap(): + """A dict nested past DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER is + replaced by the REDACTED marker instead of coming back unmasked, while the + strings sitting exactly at the cap still get the normal per-key treatment: + a sensitive one is masked and a plain one survives verbatim.""" + 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} + + 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(): + """Two models of the same shape dump into temporaries whose ids CPython + reuses as soon as the first is freed, so an id-keyed memo that does not + pin what it keys hands the second model the first one's masked copy.""" + 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 +469,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 From 093fb78bafb3c2ef8273e06370f2b072c50341d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:58:47 -0700 Subject: [PATCH 2/3] fix(masker): cut cycles at the first back-edge and walk pydantic dumps without self-recursion --- .../sensitive_data_masker.py | 6 +- .../test_sensitive_data_masker.py | 55 +++++++++++++------ 2 files changed, 41 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index b68c97e18c9..b7bd0a1498b 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -179,7 +179,8 @@ def mask_credentials_in_payload(data: object) -> object: 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. A container nested past + 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. @@ -204,6 +205,7 @@ class _PayloadWalker: 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 @@ -212,7 +214,7 @@ class _PayloadWalker: self, node: Mapping[str, object] | Sequence[object] | BaseModel, key_is_sensitive: bool, depth: int ) -> object: if isinstance(node, BaseModel): - return self._rebuild(node.model_dump(), key_is_sensitive, depth) + 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): diff --git a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py index de511b0ce11..fcdf7fb4798 100644 --- a/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py +++ b/tests/test_litellm/litellm_core_utils/test_sensitive_data_masker.py @@ -318,9 +318,6 @@ def _nested_under_levels(leaf: object, levels: int) -> object: def test_mask_credentials_in_payload_keeps_a_shared_dict_shared(): - """One dict referenced twice comes back as one masked dict referenced - twice. Rebuilding each reference separately is what turned an aliased - retry breadcrumb graph exponential in the v1.100.0 OOM.""" from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload shared: Final = {"api_key": "sk-shared-1234567890abcdef", "model": "gpt-4o-mini"} @@ -332,8 +329,6 @@ def test_mask_credentials_in_payload_keeps_a_shared_dict_shared(): def test_mask_credentials_in_payload_walks_each_dag_node_once(): - """A DAG of 9 dicts where every level references the level below three - times stays 9 dicts after masking, instead of fanning out to 3**8.""" from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload root: Final = reduce( @@ -347,10 +342,21 @@ def test_mask_credentials_in_payload_walks_each_dag_node_once(): 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(): - """The same list reached under a plain key and under a sensitive key is - masked in the sensitive spot only, whichever reference the walk meets - first, so the memo can neither leak a secret nor mask a plain value.""" from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload shared: Final = ["sk-list-1234567890abcdef"] @@ -365,9 +371,6 @@ def test_mask_credentials_in_payload_masks_a_shared_list_only_under_a_sensitive_ def test_mask_credentials_in_payload_masks_a_shared_root_model_list_only_under_a_sensitive_key(): - """A pydantic model that dumps to a list is a list once walked, so the - memo must keep its plain and sensitive rebuilds apart the same way, or - the reference met first decides what the other one shows.""" from pydantic import RootModel from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload @@ -383,11 +386,21 @@ def test_mask_credentials_in_payload_masks_a_shared_root_model_list_only_under_a 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(): - """A dict nested past DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER is - replaced by the REDACTED marker instead of coming back unmasked, while the - strings sitting exactly at the cap still get the normal per-key treatment: - a sensitive one is masked and a plain one survives verbatim.""" 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 @@ -401,6 +414,14 @@ def test_mask_credentials_in_payload_hides_containers_past_the_depth_cap(): 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), @@ -412,9 +433,7 @@ def test_mask_credentials_in_payload_hides_containers_past_the_depth_cap(): def test_mask_credentials_in_payload_keeps_sibling_models_apart(): - """Two models of the same shape dump into temporaries whose ids CPython - reuses as soon as the first is freed, so an id-keyed memo that does not - pin what it keys hands the second model the first one's masked copy.""" + """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 From aceae8e566913faf931a1c13dbbd32698695a20b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 04:07:07 -0700 Subject: [PATCH 3/3] test: drop the recursive detector allowlist entry for the removed _walk_payload --- tests/code_coverage_tests/recursive_detector.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index d8e318c61af..3c6a6a58820 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/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.