fix(sso): fall back to id_token claims when the generic userinfo response is not JSON (#46187)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/management_endpoints/ui_sso.py

@@ -989,6 +989,67 @@ def _merge_sso_token_claims(
     )
 
 
+def _id_token_claims_for_unreadable_userinfo(
+    id_token: str | None, userinfo_endpoint: str | None, error: json.JSONDecodeError
+) -> Mapping[str, object]:
+    claims: Final = _decode_sso_token_claims(id_token)
+    if not claims:
+        raise ProxyException(
+            message=(
+                f"SSO user info unavailable: userinfo endpoint {userinfo_endpoint} returned a non-JSON "
+                "response and the token response had no decodable id_token"
+            ),
+            type=ProxyErrorTypes.auth_error,
+            param="userinfo",
+            code=status.HTTP_401_UNAUTHORIZED,
+        )
+    verbose_proxy_logger.warning(
+        "Userinfo endpoint %s returned a non-JSON response (%s), falling back to id_token claims",
+        userinfo_endpoint,
+        error,
+    )
+    return claims
+
+
+class _UserinfoFallbackCapableSSO(Protocol):
+    @property
+    def access_token(self) -> str | None: ...
+
+    @property
+    def id_token(self) -> str | None: ...
+
+    async def verify_and_process(
+        self,
+        request: Request,
+        *,
+        params: dict[str, object] | None = None,  # mutable-ok: fastapi-sso verify_and_process requires dict params
+        headers: dict[str, str] | None = None,  # mutable-ok: fastapi-sso verify_and_process requires dict headers
+    ) -> CustomOpenID | None: ...
+
+
+async def _verify_generic_sso_or_fall_back_to_id_token(
+    generic_sso: _UserinfoFallbackCapableSSO,
+    request: Request,
+    token_exchange_params: dict[str, object],  # mutable-ok: fastapi-sso verify_and_process requires dict params
+    additional_headers: dict[str, str],  # mutable-ok: fastapi-sso verify_and_process requires dict headers
+    userinfo_endpoint: str | None,
+    response_convertor: Callable[[Mapping[str, object], object], CustomOpenID],
+) -> CustomOpenID | None:
+    try:
+        return await generic_sso.verify_and_process(
+            request,
+            params=token_exchange_params,
+            headers=additional_headers,
+        )
+    except json.JSONDecodeError as userinfo_error:
+        if generic_sso.access_token is None:
+            raise
+        return response_convertor(
+            _id_token_claims_for_unreadable_userinfo(generic_sso.id_token, userinfo_endpoint, userinfo_error),
+            generic_sso,
+        )
+
+
 async def _raise_if_sso_exceeds_free_user_limit(premium_user: bool, prisma_client: PrismaClient | None) -> None:
     """Free tier allows SSO for up to 5 billable users; beyond that requires an Enterprise license."""
     if premium_user is True:
@@ -1655,10 +1716,15 @@ async def get_generic_sso_response(
             # process_sso_jwt_access_token can extract JWT-embedded roles/teams.
             access_token_str: str | None = combined_response.get("access_token")
         else:
-            result = await generic_sso.verify_and_process(
-                request,
-                params=token_exchange_params,
-                headers=additional_generic_sso_headers_dict,
+            result = await _verify_generic_sso_or_fall_back_to_id_token(
+                generic_sso=cast(  # cast-ok: fastapi-sso builds the provider class dynamically via create_provider
+                    _UserinfoFallbackCapableSSO, generic_sso
+                ),
+                request=request,
+                token_exchange_params=token_exchange_params,
+                additional_headers=additional_generic_sso_headers_dict,
+                userinfo_endpoint=generic_userinfo_endpoint,
+                response_convertor=response_convertor,
             )
             access_token_str = generic_sso.access_token
             sso_assertion = assertion_from_sso_login(generic_sso.id_token, generic_sso.refresh_token)

tests/unit/proxy/management_endpoints/test_ui_sso.py

@@ -8914,6 +8914,143 @@ async def test_verify_and_process_arm_captures_sso_assertion():
     assert sso_assertion.refresh_token.get_secret_value() == "rt_from_idp"
 
 
+_GENERIC_SSO_CALLBACK_ENV = {
+    "GENERIC_CLIENT_SECRET": "test_secret",
+    "GENERIC_AUTHORIZATION_ENDPOINT": "https://auth.example.com/auth",
+    "GENERIC_TOKEN_ENDPOINT": "https://auth.example.com/token",
+    "GENERIC_USERINFO_ENDPOINT": "https://auth.example.com/userinfo",
+    "GENERIC_CLIENT_USE_PKCE": "",
+}
+
+
+def _generic_sso_callback_request() -> Request:
+    from starlette.requests import Request as StarletteRequest
+
+    return StarletteRequest(
+        {
+            "type": "http",
+            "method": "GET",
+            "scheme": "https",
+            "server": ("proxy.example.com", 443),
+            "path": "/sso/callback",
+            "query_string": b"code=abc",
+            "headers": [],
+        }
+    )
+
+
+def _generic_sso_jwt_handler() -> MagicMock:
+    jwt_handler = MagicMock(spec=JWTHandler)
+    jwt_handler.get_team_ids_from_jwt.return_value = []
+    return jwt_handler
+
+
+@pytest.mark.asyncio
+async def test_non_json_userinfo_falls_back_to_id_token_claims():
+    """A userinfo endpoint answering with a non-JSON body (ADFS 401, Content-Length 0)
+    must fall back to id_token claims on the non-PKCE arm instead of surfacing a 500."""
+    import jwt as pyjwt
+
+    from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response
+
+    id_token = pyjwt.encode(
+        {"sub": "u-adfs", "preferred_username": "u-adfs", "email": "adfs@example.com"},
+        "test-idp-signing-key-32-bytes-long-xxxx",
+        algorithm="HS256",
+    )
+    request = _generic_sso_callback_request()
+
+    with (
+        respx.mock,
+        patch.dict(os.environ, _GENERIC_SSO_CALLBACK_ENV),
+    ):
+        respx.post("https://auth.example.com/token").mock(
+            return_value=httpx.Response(200, json={"access_token": "at", "token_type": "Bearer", "id_token": id_token})
+        )
+        userinfo_route = respx.get("https://auth.example.com/userinfo").mock(
+            return_value=httpx.Response(401, content=b"", headers={"content-type": "text/html"})
+        )
+        result, received_response, _, _ = await get_generic_sso_response(
+            request=request,
+            jwt_handler=_generic_sso_jwt_handler(),
+            generic_client_id="test_client_id",
+            redirect_url="https://proxy.example.com/sso/callback",
+            sso_jwt_handler=None,
+        )
+
+    assert result.email == "adfs@example.com"
+    assert result.id == "u-adfs"
+    assert received_response is not None
+    assert received_response["email"] == "adfs@example.com"
+    assert received_response["sub"] == "u-adfs"
+    assert "id_token" not in received_response
+    assert "access_token" not in received_response
+    assert userinfo_route.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_non_json_userinfo_without_id_token_raises_401():
+    """Same unreadable userinfo but no decodable id_token in the token response:
+    the login must fail closed with a 401 ProxyException naming the userinfo endpoint."""
+    from litellm.proxy._types import ProxyException
+    from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response
+
+    request = _generic_sso_callback_request()
+
+    with (
+        respx.mock,
+        patch.dict(os.environ, _GENERIC_SSO_CALLBACK_ENV),
+    ):
+        respx.post("https://auth.example.com/token").mock(
+            return_value=httpx.Response(200, json={"access_token": "at", "token_type": "Bearer"})
+        )
+        respx.get("https://auth.example.com/userinfo").mock(
+            return_value=httpx.Response(401, content=b"", headers={"content-type": "text/html"})
+        )
+        with pytest.raises(ProxyException) as exc_info:
+            await get_generic_sso_response(
+                request=request,
+                jwt_handler=_generic_sso_jwt_handler(),
+                generic_client_id="test_client_id",
+                redirect_url="https://proxy.example.com/sso/callback",
+                sso_jwt_handler=None,
+            )
+
+    assert str(exc_info.value.code) == "401"
+    assert "https://auth.example.com/userinfo" in exc_info.value.message
+
+
+@pytest.mark.asyncio
+async def test_non_json_token_response_is_not_swallowed_by_userinfo_fallback():
+    """A non-JSON TOKEN endpoint response must keep the old behaviour (raw JSONDecodeError
+    through _handle_generic_sso_error): no access token was parsed, so the id_token
+    fallback must not fire and userinfo must never be called."""
+    from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response
+
+    request = _generic_sso_callback_request()
+
+    with (
+        respx.mock,
+        patch.dict(os.environ, _GENERIC_SSO_CALLBACK_ENV),
+    ):
+        respx.post("https://auth.example.com/token").mock(
+            return_value=httpx.Response(200, content=b"", headers={"content-type": "text/html"})
+        )
+        userinfo_route = respx.get("https://auth.example.com/userinfo").mock(
+            return_value=httpx.Response(200, json={"sub": "unused"})
+        )
+        with pytest.raises(json.JSONDecodeError):
+            await get_generic_sso_response(
+                request=request,
+                jwt_handler=_generic_sso_jwt_handler(),
+                generic_client_id="test_client_id",
+                redirect_url="https://proxy.example.com/sso/callback",
+                sso_jwt_handler=None,
+            )
+
+    assert not userinfo_route.called
+
+
 @pytest.mark.asyncio
 async def test_redirect_from_openid_persists_assertion_under_canonical_user_id():
     """The browser funnel persists the captured assertion AFTER canonical user