address greptile review feedback (greploop iteration 4)

litellm/proxy/management_endpoints/ui_sso.py

@@ -2492,11 +2492,12 @@ class SSOAuthenticationHandler:
         token_params: Dict[str, Any] = {"include_client_id": generic_include_client_id}
 
         # Retrieve PKCE code_verifier if PKCE was used in authorization.
-        # Use same cache as store: Redis when available (multi-pod), else in-memory.
+        # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip
+        # on every non-PKCE SSO callback.
         query_params = dict(request.query_params)
         state = query_params.get("state")
 
-        if state:
+        if state and os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true":
             from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache
 
             cache_key = f"pkce_verifier:{state}"

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -4,6 +4,7 @@ import os
 import sys
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
 from fastapi import Request
 
@@ -3147,7 +3148,7 @@ class TestPKCEFunctionality:
         mock_cache.async_get_cache = AsyncMock(return_value=test_code_verifier)
         mock_cache.async_delete_cache = AsyncMock()
 
-        with patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache):
+        with patch("litellm.proxy.proxy_server.redis_usage_cache", None), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}):
             # Act
             token_params = (
                 await SSOAuthenticationHandler.prepare_token_exchange_parameters(
@@ -4503,12 +4504,6 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch):
 # Tests for SSOAuthenticationHandler PKCE token exchange methods
 # ──────────────────────────────────────────────────────────────────────────────
 
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import httpx
-import pytest
-
-
 @pytest.mark.asyncio
 async def test_pkce_token_exchange_basic_auth():
     """When include_client_id=False, client credentials go via HTTP Basic Auth."""
@@ -4534,9 +4529,6 @@ async def test_pkce_token_exchange_basic_auth():
         assert isinstance(kwargs["auth"], httpx.BasicAuth)
         return mock_response
 
-    async def fake_get(*args, **kwargs):
-        return mock_userinfo_response
-
     with patch("httpx.AsyncClient") as mock_client_cls:
         mock_client = AsyncMock()
         mock_client.__aenter__ = AsyncMock(return_value=mock_client)