address greptile review feedback (greploop iteration 32)

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(

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",