From ee703cea99252c0e45b2a74b9f376e2ecda9ea65 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 27 Feb 2026 20:30:47 -0800 Subject: [PATCH] fix(jwt): OIDC discovery URLs, roles array handling, dot-notation error hints (#22336) * fix(jwt): support OIDC discovery URLs, handle roles array, improve error hints Three fixes for Azure AD JWT auth: 1. OIDC discovery URL support - JWT_PUBLIC_KEY_URL can now be set to .well-known/openid-configuration endpoints. The proxy fetches the discovery doc, extracts jwks_uri, and caches it. 2. Handle roles claim as array - when team_id_jwt_field points to a list (e.g. AAD's "roles": ["team1"]), auto-unwrap the first element instead of crashing with 'unhashable type: list'. 3. Better error hint for dot-notation indexing - when team_id_jwt_field is set to "roles.0" or "roles[0]", the 401 error now explains to use "roles" instead and that LiteLLM auto-unwraps lists. * Add integration demo script for JWT auth fixes (OIDC discovery, array roles, dot-notation hints) Co-authored-by: Ishaan Jaff * Add demo_servers.py for manual JWT auth testing with mock JWKS/OIDC endpoints Co-authored-by: Ishaan Jaff * Add demo screenshots for PR comment Co-authored-by: Ishaan Jaff * Add integration test results with screenshots for PR review Co-authored-by: Ishaan Jaff * address greptile review feedback (greploop iteration 1) - fix: add HTTP status code check in _resolve_jwks_url before parsing JSON - fix: remove misleading bracket-notation hint from debug log (get_nested_value does not support it) * Update tests/test_litellm/proxy/auth/test_handle_jwt.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * remove demo scripts and assets --------- Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/auth/handle_jwt.py | 84 +++++- .../proxy/auth/test_handle_jwt.py | 269 ++++++++++++++++++ 2 files changed, 351 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9921b74b56..553ba4d6c4 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -8,6 +8,7 @@ JWT token must have 'litellm_proxy_admin' in scope. import fnmatch import os +import re from typing import Any, List, Literal, Optional, Set, Tuple, cast from cryptography import x509 @@ -235,7 +236,17 @@ class JWTHandler: return self.litellm_jwtauth.team_id_default else: return default_value - # At this point, team_id is not the sentinel, so it should be a string + # AAD and other IdPs often send roles/groups as a list of strings. + # team_id_jwt_field is singular, so take the first element when a list + # is returned. This avoids "unhashable type: 'list'" errors downstream. + if isinstance(team_id, list): + if not team_id: + return default_value + verbose_proxy_logger.debug( + f"JWT Auth: team_id_jwt_field '{self.litellm_jwtauth.team_id_jwt_field}' " + f"returned a list {team_id}; using first element '{team_id[0]}' automatically." + ) + team_id = team_id[0] return team_id # type: ignore[return-value] elif self.litellm_jwtauth.team_id_default is not None: team_id = self.litellm_jwtauth.team_id_default @@ -453,6 +464,52 @@ class JWTHandler: scopes = [] return scopes + async def _resolve_jwks_url(self, url: str) -> str: + """ + If url points to an OIDC discovery document (*.well-known/openid-configuration), + fetch it and return the jwks_uri contained within. Otherwise return url unchanged. + This lets JWT_PUBLIC_KEY_URL be set to a well-known discovery endpoint instead of + requiring operators to manually find the JWKS URL. + """ + if ".well-known/openid-configuration" not in url: + return url + + cache_key = f"litellm_oidc_discovery_{url}" + cached_jwks_uri = await self.user_api_key_cache.async_get_cache(cache_key) + if cached_jwks_uri is not None: + return cached_jwks_uri + + verbose_proxy_logger.debug( + f"JWT Auth: Fetching OIDC discovery document from {url}" + ) + response = await self.http_handler.get(url) + if response.status_code != 200: + raise Exception( + f"JWT Auth: OIDC discovery endpoint {url} returned status {response.status_code}: {response.text}" + ) + try: + discovery = response.json() + except Exception as e: + raise Exception( + f"JWT Auth: Failed to parse OIDC discovery document at {url}: {e}" + ) + + jwks_uri = discovery.get("jwks_uri") + if not jwks_uri: + raise Exception( + f"JWT Auth: OIDC discovery document at {url} does not contain a 'jwks_uri' field." + ) + + verbose_proxy_logger.debug( + f"JWT Auth: Resolved OIDC discovery {url} -> jwks_uri={jwks_uri}" + ) + await self.user_api_key_cache.async_set_cache( + key=cache_key, + value=jwks_uri, + ttl=self.litellm_jwtauth.public_key_ttl, + ) + return jwks_uri + async def get_public_key(self, kid: Optional[str]) -> dict: keys_url = os.getenv("JWT_PUBLIC_KEY_URL") @@ -462,6 +519,7 @@ class JWTHandler: keys_url_list = [url.strip() for url in keys_url.split(",")] for key_url in keys_url_list: + key_url = await self._resolve_jwks_url(key_url) cache_key = f"litellm_jwt_auth_keys_{key_url}" cached_keys = await self.user_api_key_cache.async_get_cache(cache_key) @@ -913,8 +971,30 @@ class JWTAuthManager: if jwt_handler.is_required_team_id() is True: team_id_field = jwt_handler.litellm_jwtauth.team_id_jwt_field team_alias_field = jwt_handler.litellm_jwtauth.team_alias_jwt_field + hint = "" + if team_id_field: + # "roles.0" — dot-notation numeric indexing is not supported + if "." in team_id_field: + parts = team_id_field.rsplit(".", 1) + if parts[-1].isdigit(): + base_field = parts[0] + hint = ( + f" Hint: dot-notation array indexing (e.g. '{team_id_field}') is not " + f"supported. Use '{base_field}' instead — LiteLLM automatically " + f"uses the first element when the field value is a list." + ) + # "roles[0]" — bracket-notation indexing is also not supported in get_nested_value + elif "[" in team_id_field and team_id_field.endswith("]"): + m = re.match(r"^(\w+)\[(\d+)\]$", team_id_field) + if m: + base_field = m.group(1) + hint = ( + f" Hint: array indexing (e.g. '{team_id_field}') is not supported " + f"in team_id_jwt_field. Use '{base_field}' instead — LiteLLM " + f"automatically uses the first element when the field value is a list." + ) raise Exception( - f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'" + f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'.{hint}" ) return individual_team_id, team_object diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index b56d13bb93..8418dde5e9 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1485,4 +1485,273 @@ async def test_get_objects_resolves_org_by_name(): ) +# --------------------------------------------------------------------------- +# Fix 1: OIDC discovery URL resolution +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_passthrough_for_direct_jwks_url(): + """Non-discovery URLs are returned unchanged.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + url = "https://login.microsoftonline.com/common/discovery/keys" + result = await handler._resolve_jwks_url(url) + assert result == url + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_resolves_oidc_discovery_document(): + """ + A .well-known/openid-configuration URL should be fetched and its + jwks_uri returned. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + cache = DualCache() + handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"jwks_uri": jwks_url, "issuer": "https://..."} + + with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get: + result = await handler._resolve_jwks_url(discovery_url) + + assert result == jwks_url + mock_get.assert_called_once_with(discovery_url) + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_caches_resolved_jwks_uri(): + """Resolved jwks_uri is cached — second call does not hit the network.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + cache = DualCache() + handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + discovery_url = "https://login.microsoftonline.com/tenant/.well-known/openid-configuration" + jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" + + mock_response = MagicMock() + mock_response.json.return_value = {"jwks_uri": jwks_url} + + with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get: + first = await handler._resolve_jwks_url(discovery_url) + second = await handler._resolve_jwks_url(discovery_url) + + assert first == jwks_url + assert second == jwks_url + # Network should only be hit once + assert mock_get.call_count == 1 + + +@pytest.mark.asyncio +async def test_resolve_jwks_url_raises_if_no_jwks_uri_in_discovery_doc(): + """Raise a helpful error if the discovery document has no jwks_uri.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + discovery_url = "https://example.com/.well-known/openid-configuration" + mock_response = MagicMock() + mock_response.json.return_value = {"issuer": "https://example.com"} # no jwks_uri + + with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response): + with pytest.raises(Exception, match="jwks_uri"): + await handler._resolve_jwks_url(discovery_url) + + +# --------------------------------------------------------------------------- +# Fix 2: handle array values in team_id_jwt_field (e.g. AAD "roles" claim) +# --------------------------------------------------------------------------- + + +def _make_jwt_handler(team_id_jwt_field: str) -> JWTHandler: + from litellm.caching.dual_cache import DualCache + + handler = JWTHandler() + handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(team_id_jwt_field=team_id_jwt_field), + ) + return handler + + +def test_get_team_id_returns_first_element_when_roles_is_list(): + """ + AAD sends roles as a list. get_team_id() must return the first string + element rather than the raw list (which would later crash with + 'unhashable type: list'). + """ + handler = _make_jwt_handler("roles") + token = {"oid": "user-oid", "roles": ["team1"]} + result = handler.get_team_id(token=token, default_value=None) + assert result == "team1" + + +def test_get_team_id_returns_first_element_from_multi_value_roles_list(): + """When roles has multiple entries, the first one is used.""" + handler = _make_jwt_handler("roles") + token = {"roles": ["team2", "team1"]} + result = handler.get_team_id(token=token, default_value=None) + assert result == "team2" + + +def test_get_team_id_returns_default_when_roles_list_is_empty(): + """Empty list should fall back to default_value.""" + handler = _make_jwt_handler("roles") + token = {"roles": []} + result = handler.get_team_id(token=token, default_value="fallback") + assert result == "fallback" + + +def test_get_team_id_still_works_with_string_value(): + """String values (non-array) continue to work as before.""" + handler = _make_jwt_handler("appid") + token = {"appid": "my-team-id"} + result = handler.get_team_id(token=token, default_value=None) + assert result == "my-team-id" + + +def test_get_team_id_list_result_is_hashable(): + """ + The value returned by get_team_id() must be hashable so it can be + added to a set (the operation that previously crashed). + """ + handler = _make_jwt_handler("roles") + token = {"roles": ["team1"]} + result = handler.get_team_id(token=token, default_value=None) + # This must not raise TypeError + s: set = set() + s.add(result) + assert "team1" in s + + +# --------------------------------------------------------------------------- +# Fix 3: helpful error message for dot-notation array indexing (roles.0) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_hints_bracket_notation(): + """ + When team_id_jwt_field is set to 'roles.0' (unsupported dot-notation for + array indexing) and no team is found, the exception message should suggest + using 'roles' instead (and explain LiteLLM auto-unwraps list values). + """ + from unittest.mock import MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = _make_jwt_handler("roles.0") + # token has roles as a list — dot-notation won't find anything + token = {"roles": ["team1"]} + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + error_msg = str(exc_info.value) + # Should mention the bad field name and suggest the fix + assert "roles.0" in error_msg, f"Expected field name in: {error_msg}" + assert "roles" in error_msg and "list" in error_msg, ( + f"Expected hint about using 'roles' instead: {error_msg}" + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_hints_bracket_index_notation(): + """ + When team_id_jwt_field is set to 'roles[0]' (bracket indexing, also unsupported + in get_nested_value) the error message should suggest using 'roles' instead. + """ + from unittest.mock import MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = _make_jwt_handler("roles[0]") + token = {"roles": ["team1"]} + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + error_msg = str(exc_info.value) + assert "roles[0]" in error_msg, f"Expected field name in: {error_msg}" + assert "roles" in error_msg and "list" in error_msg, ( + f"Expected hint about using 'roles' instead: {error_msg}" + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_no_hint_for_valid_field(): + """ + When team_id_jwt_field is a normal field name (no dot-notation) the + error message should not contain a spurious bracket-notation hint. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.dual_cache import DualCache + + handler = _make_jwt_handler("appid") + token = {} # no appid — triggers the "no team found" path + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=DualCache(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + error_msg = str(exc_info.value) + assert "Hint" not in error_msg