diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index daf9e2e44ab..c77bb9b5b44 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -2748,12 +2748,14 @@ class SSOAuthenticationHandler: if client_secret: token_data["client_secret"] = client_secret - # The try/except is INSIDE the async with so it only covers the POST call. - # If TLS teardown in __aexit__ raises, it propagates after the try/except - # has already completed — it is NOT caught here and NOT mis-labelled as - # "Token endpoint request failed". httpx buffers the full response body - # inside the context, so status_code / text / json() are safe to access - # outside the async with block. + # Initialize response to None — guards against an UnboundLocalError in the + # unlikely case where httpx.AsyncClient() construction itself raises before + # the POST is attempted. The try/except is INSIDE the async with so that + # TLS teardown exceptions from __aexit__ propagate as-is and are NOT + # mis-labelled as "Token endpoint request failed". httpx buffers the full + # response body before __aexit__, so status_code / text / json() remain + # valid after the context exits. + response = None async with httpx.AsyncClient() as http_client: try: response = await http_client.post(token_endpoint, **post_kwargs) @@ -2769,6 +2771,15 @@ class SSOAuthenticationHandler: code=status.HTTP_401_UNAUTHORIZED, ) from exc + if response is None: + # Should never happen in practice — construction failure is unexpected. + raise ProxyException( + message="Token endpoint request did not return a response", + type=ProxyErrorTypes.auth_error, + param="token_exchange", + code=status.HTTP_401_UNAUTHORIZED, + ) + # Response processing outside the async with — httpx buffers the full # response body so status_code / text / json() remain valid after __aexit__. if response.status_code != 200: @@ -2914,6 +2925,14 @@ class SSOAuthenticationHandler: if userinfo is None and isinstance(id_token, str) and id_token: try: userinfo = jwt.decode(id_token, options={"verify_signature": False}) + if not userinfo: + # jwt.decode returned an empty dict (payload-free JWT or provider bug). + # Treat this the same as a missing userinfo — the session would have no + # identity claims, which is equivalent to a broken session. + verbose_proxy_logger.warning( + "id_token decoded to an empty payload — treating as failure." + ) + userinfo = None except Exception as decode_err: verbose_proxy_logger.error("Failed to decode id_token: %s", decode_err) raise ProxyException( diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 5dfbf545860..22122daeee6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3424,13 +3424,20 @@ class TestPKCEFunctionality: assert kwargs.get("data", {}).get("code_verifier") == "verifier_abc" return mock_response + # Use separate mock clients for token exchange and userinfo — + # each httpx.AsyncClient() call gets its own independent mock. + mock_token_client = AsyncMock() + mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client) + mock_token_client.__aexit__ = AsyncMock(return_value=False) + mock_token_client.post = AsyncMock(side_effect=fake_post) + + mock_userinfo_client = AsyncMock() + mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client) + mock_userinfo_client.__aexit__ = AsyncMock(return_value=False) + mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo_response) + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: - mock_client = AsyncMock() - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.post = AsyncMock(side_effect=fake_post) - mock_client.get = AsyncMock(return_value=mock_userinfo_response) - mock_client_cls.return_value = mock_client + mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_code_123", @@ -3447,7 +3454,7 @@ class TestPKCEFunctionality: assert result["access_token"] == "tok_abc" assert result["email"] == "user@example.com" # Verify userinfo GET used the correct Bearer token header - get_call = mock_client.get.call_args + get_call = mock_userinfo_client.get.call_args assert get_call is not None assert get_call.kwargs["headers"]["Authorization"] == "Bearer tok_abc" @@ -3474,16 +3481,22 @@ class TestPKCEFunctionality: mock.json.return_value = token_resp return mock + mock_userinfo = MagicMock() + mock_userinfo.status_code = 200 + mock_userinfo.json.return_value = userinfo_resp + + mock_token_client = AsyncMock() + mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client) + mock_token_client.__aexit__ = AsyncMock(return_value=False) + mock_token_client.post = AsyncMock(side_effect=fake_post) + + mock_userinfo_client = AsyncMock() + mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client) + mock_userinfo_client.__aexit__ = AsyncMock(return_value=False) + mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: - mock_client = AsyncMock() - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.post = AsyncMock(side_effect=fake_post) - mock_userinfo = MagicMock() - mock_userinfo.status_code = 200 - mock_userinfo.json.return_value = userinfo_resp - mock_client.get = AsyncMock(return_value=mock_userinfo) - mock_client_cls.return_value = mock_client + mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_code_456", @@ -3675,16 +3688,22 @@ class TestPKCEFunctionality: mock.json.return_value = token_resp return mock + mock_userinfo = MagicMock() + mock_userinfo.status_code = 200 + mock_userinfo.json.return_value = userinfo_resp + + mock_token_client = AsyncMock() + mock_token_client.__aenter__ = AsyncMock(return_value=mock_token_client) + mock_token_client.__aexit__ = AsyncMock(return_value=False) + mock_token_client.post = AsyncMock(side_effect=fake_post) + + mock_userinfo_client = AsyncMock() + mock_userinfo_client.__aenter__ = AsyncMock(return_value=mock_userinfo_client) + mock_userinfo_client.__aexit__ = AsyncMock(return_value=False) + mock_userinfo_client.get = AsyncMock(return_value=mock_userinfo) + with patch("litellm.proxy.management_endpoints.ui_sso.httpx.AsyncClient") as mock_client_cls: - mock_client = AsyncMock() - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=False) - mock_client.post = AsyncMock(side_effect=fake_post) - mock_userinfo = MagicMock() - mock_userinfo.status_code = 200 - mock_userinfo.json.return_value = userinfo_resp - mock_client.get = AsyncMock(return_value=mock_userinfo) - mock_client_cls.return_value = mock_client + mock_client_cls.side_effect = [mock_token_client, mock_userinfo_client] result = await SSOAuthenticationHandler._pkce_token_exchange( authorization_code="auth_pub",