From 8a667f096db228420f000bb3e0d3810ac24c6ec0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 5 Mar 2026 14:02:55 -0800 Subject: [PATCH] address greptile review feedback (greploop iteration 4) --- litellm/proxy/management_endpoints/ui_sso.py | 5 +++-- .../proxy/management_endpoints/test_ui_sso.py | 12 ++---------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 8ecb6b4a629..a5b1f3d9978 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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}" 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 08ea3cf000e..7d1ba380495 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/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)