Merge 784adbc6ab into a2bf67a037

litellm/litellm_core_utils/dot_notation_indexing.py

@@ -28,6 +28,9 @@ from typing import Any, Final, TypeVar
 
 T = TypeVar("T")
 
+_MISSING: Final = object()
+"""Sentinel for "this path does not resolve", distinct from any caller default."""
+
 
 def get_nested_value(data: Mapping[str, object], key_path: str, default: T | None = None) -> T | None:
     """
@@ -54,29 +57,45 @@ def get_nested_value(data: Mapping[str, object], key_path: str, default: T | Non
     if not key_path:
         return default
 
-    # Remove metadata. prefix if it exists
-    key_path = key_path.replace("metadata.", "", 1) if key_path.startswith("metadata.") else key_path
-
     # Split the key path into parts, respecting escaped dots (\.)
     # Use a temporary placeholder, split on unescaped dots, then restore
     placeholder: Final = "\x00"
-    parts = key_path.replace("\\.", placeholder).split(".")
-    parts = [p.replace(placeholder, ".") for p in parts]
 
-    # Traverse through the dictionary
-    current: Any = data
-    for part in parts:
-        try:
-            current = current[part]
-        except (KeyError, TypeError):
-            return default
+    def _resolve(path: str) -> Any:
+        parts = path.replace("\\.", placeholder).split(".")
+        current: Any = data
+        for part in parts:
+            part = part.replace(placeholder, ".")
+            try:
+                current = current[part]
+            except (KeyError, TypeError):
+                return _MISSING
+        return current
+
+    found = _resolve(key_path)
+
+    # Historically a leading "metadata." was stripped unconditionally, so that
+    # `roles_jwt_field: "metadata.roles"` also matched a token carrying `roles`
+    # at the top level. Stripping it up front also made a genuinely nested
+    # `metadata` object unreachable: the path became `roles`, which is absent,
+    # so the claim read as missing and JWT auth fell back to its defaults --
+    # get_team_id() then places the caller in `team_id_default` rather than the
+    # team they named. Try the path as written first and keep the legacy
+    # spelling working as a fallback, so both token shapes resolve.
+    if found is _MISSING and key_path.startswith("metadata."):
+        without_prefix: Final = key_path[len("metadata.") :]
+        if without_prefix:
+            found = _resolve(without_prefix)
+
+    if found is _MISSING:
+        return default
 
     # If default is None, we can return any type
     if default is None:
-        return current
+        return found
 
     # Otherwise, ensure the type matches the default
-    return current if isinstance(current, type(default)) else default
+    return found if isinstance(found, type(default)) else default
 
 
 def _parse_path_segments(path: str) -> list:

tests/unit/litellm_core_utils/test_dot_notation_indexing.py

@@ -41,6 +41,53 @@ class TestGetNestedValue:
         data = {"user": {"email": "test@example.com"}}
         assert get_nested_value(data, "metadata.user.email") == "test@example.com"
 
+    def test_nested_metadata_object_is_reachable(self):
+        """A token that really nests claims under `metadata` resolves.
+
+        The prefix used to be stripped before traversal, so `metadata.roles`
+        was looked up as `roles` and a genuinely nested claim read as missing.
+        JWT auth then fell back to its defaults, which for get_team_id means
+        `team_id_default` rather than the team the token named.
+        """
+        data = {"metadata": {"roles": ["admin"]}}
+        assert get_nested_value(data, "metadata.roles") == ["admin"]
+
+    def test_nested_metadata_wins_over_top_level(self):
+        """When both shapes exist, the literal path is the one that was asked for."""
+        data = {"metadata": {"roles": ["nested"]}, "roles": ["top-level"]}
+        assert get_nested_value(data, "metadata.roles") == ["nested"]
+
+    def test_metadata_prefix_still_falls_back_to_top_level(self):
+        """The legacy spelling keeps working when there is no nested object."""
+        data = {"roles": ["admin"]}
+        assert get_nested_value(data, "metadata.roles") == ["admin"]
+
+    def test_metadata_prefix_falls_back_to_a_default(self):
+        """With neither shape present the caller's default comes back unchanged."""
+        data = {"other": 1}
+        assert get_nested_value(data, "metadata.roles", "fallback") == "fallback"
+
+    def test_metadata_prefix_applies_only_once(self):
+        """Only a leading `metadata.` is treated as a prefix.
+
+        `metadata.metadata.x` strips one segment, so a key literally named
+        `metadata` inside a `metadata` object is still reachable.
+        """
+        data = {"metadata": {"metadata": {"team_id": "team-1"}}}
+        assert get_nested_value(data, "metadata.metadata.team_id") == "team-1"
+
+    def test_nested_metadata_respects_the_default_type(self):
+        """The caller's default still governs type coercion on the nested read."""
+        data = {"metadata": {"team_id": "team-1"}}
+        assert get_nested_value(data, "metadata.team_id", "fallback") == "team-1"
+        # A default of another type means the nested value is not usable here,
+        # exactly as for a non-prefixed path.
+        assert get_nested_value(data, "metadata.team_id", 1234) == 1234
+
+    def test_bare_metadata_key_is_untouched(self):
+        """A top-level key named `metadata` resolves as written."""
+        assert get_nested_value({"metadata": "value"}, "metadata") == "value"
+
     def test_escaped_dot_in_key(self):
         """Test accessing keys that contain dots using escape sequence."""
         data = {"kubernetes.io": {"namespace": "default"}}