From da8f88bebaaa7f992c3ab8e2bd70faa1f2b86ee8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 16 Oct 2025 12:13:54 -0700 Subject: [PATCH] init changes to add PKCE for OKTA --- ...odel_prices_and_context_window_backup.json | 100 +++++++++++- litellm/proxy/management_endpoints/ui_sso.py | 148 +++++++++++++++++- 2 files changed, 240 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 209faf999d..caca89d23f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9604,6 +9604,54 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-2.5-flash-image": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "max_pdf_size_mb": 30, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "image_generation", + "output_cost_per_image": 0.039, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 8000000 + }, "gemini-2.5-flash-image-preview": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 1e-06, @@ -11057,6 +11105,54 @@ "supports_web_search": true, "tpm": 8000000 }, + "gemini/gemini-2.5-flash-image": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_audio_token": 1e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "max_pdf_size_mb": 30, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "image_generation", + "output_cost_per_image": 0.039, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, + "rpm": 100000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-2.5-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 8000000 + }, "gemini/gemini-2.5-flash-image-preview": { "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 1e-06, @@ -13315,11 +13411,11 @@ "text" ], "supports_function_calling": true, - "supports_native_streaming": false, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, - "supports_reasoning": false, + "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": false, "supports_tool_choice": true, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 54909b9712..a8dc88dd1f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -9,7 +9,10 @@ Has all /sso/* routes """ import asyncio +import base64 +import hashlib import os +import secrets from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast @@ -381,7 +384,10 @@ async def get_generic_sso_response( try: result = await generic_sso.verify_and_process( request, - params={"include_client_id": generic_include_client_id}, + params=SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=request, + generic_include_client_id=generic_include_client_id, + ), headers=additional_generic_sso_headers_dict, ) @@ -1072,14 +1078,66 @@ class SSOAuthenticationHandler: # or a cryptographicly signed state that we can verify stateless # For simplification we are using a static state, this is not perfect but some # SSO providers do not allow stateless verification - redirect_params = ( + redirect_params, code_verifier = ( SSOAuthenticationHandler._get_generic_sso_redirect_params( state=state, generic_authorization_endpoint=generic_authorization_endpoint, ) ) - return await generic_sso.get_login_redirect(**redirect_params) # type: ignore + # Separate PKCE params from state params (fastapi-sso doesn't accept code_challenge) + pkce_params = {} + state_only_params = {} + for key, value in redirect_params.items(): + if key in ("code_challenge", "code_challenge_method"): + pkce_params[key] = value + else: + state_only_params[key] = value + + # Get the redirect response from fastapi-sso with only state param + redirect_response = await generic_sso.get_login_redirect(**state_only_params) # type: ignore + + # If PKCE is enabled, add PKCE parameters to the redirect URL + if code_verifier and "state" in redirect_params: + from urllib.parse import parse_qs, urlencode, urlparse, urlunparse + + from litellm.proxy.proxy_server import user_api_key_cache + + # Store code_verifier in cache (10 min TTL) + cache_key = f"pkce_verifier:{redirect_params['state']}" + user_api_key_cache.set_cache( + key=cache_key, + value=code_verifier, + ttl=600, + ) + + # Add PKCE parameters to the authorization URL + if pkce_params: + parsed_url = urlparse(str(redirect_response.headers["location"])) + query_params = parse_qs(parsed_url.query) + + # Add PKCE parameters + for key, value in pkce_params.items(): + query_params[key] = [value] + + # Reconstruct the URL with PKCE parameters + new_query = urlencode(query_params, doseq=True) + new_url = urlunparse(( + parsed_url.scheme, + parsed_url.netloc, + parsed_url.path, + parsed_url.params, + new_query, + parsed_url.fragment + )) + + # Update the redirect response + redirect_response.headers["location"] = new_url + verbose_proxy_logger.debug( + "PKCE parameters added to authorization URL" + ) + + return redirect_response raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" ) @@ -1088,9 +1146,10 @@ class SSOAuthenticationHandler: def _get_generic_sso_redirect_params( state: Optional[str] = None, generic_authorization_endpoint: Optional[str] = None, - ) -> dict: + ) -> Tuple[dict, Optional[str]]: """ Get redirect parameters for Generic SSO with proper state priority handling. + Optionally generates PKCE parameters if GENERIC_CLIENT_USE_PKCE is enabled. Priority order: 1. CLI state (if provided) @@ -1102,9 +1161,12 @@ class SSOAuthenticationHandler: generic_authorization_endpoint: Authorization endpoint URL Returns: - dict: Redirect parameters for SSO login + Tuple[dict, Optional[str]]: + - Redirect parameters for SSO login (may include PKCE params) + - code_verifier (if PKCE is enabled, None otherwise) """ redirect_params = {} + code_verifier: Optional[str] = None if state: # CLI state takes priority @@ -1122,7 +1184,18 @@ class SSOAuthenticationHandler: "state" ] = uuid.uuid4().hex # set state param for okta - required - return redirect_params + # Handle PKCE (Proof Key for Code Exchange) if enabled + # Set GENERIC_CLIENT_USE_PKCE=true to enable PKCE for enhanced OAuth security + use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" + if use_pkce: + code_verifier, code_challenge = SSOAuthenticationHandler.generate_pkce_params() + redirect_params["code_challenge"] = code_challenge + redirect_params["code_challenge_method"] = "S256" + verbose_proxy_logger.debug( + "PKCE enabled - code_challenge added to authorization request" + ) + + return redirect_params, code_verifier @staticmethod def should_use_sso_handler( @@ -1606,6 +1679,69 @@ class SSOAuthenticationHandler: redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) redirect_response.set_cookie(key="token", value=jwt_token) return redirect_response + + + @staticmethod + def prepare_token_exchange_parameters( + request: Request, + generic_include_client_id: bool, + ) -> dict: + """ + Prepare token exchange parameters for Generic SSO. + + Args: + request: Request object + generic_include_client_id: Generic OAuth Client ID + + Returns: + dict: Token exchange parameters + """ + # Prepare token exchange parameters + token_params = {"include_client_id": generic_include_client_id} + + # Retrieve PKCE code_verifier if PKCE was used in authorization + query_params = dict(request.query_params) + state = query_params.get("state") + if state: + from litellm.proxy.proxy_server import user_api_key_cache + + cache_key = f"pkce_verifier:{state}" + code_verifier = user_api_key_cache.get_cache(key=cache_key) + + if code_verifier: + # Add code_verifier to token exchange parameters + token_params["code_verifier"] = code_verifier + verbose_proxy_logger.debug( + "PKCE code_verifier retrieved and will be included in token exchange" + ) + + # Clean up the cache entry (single-use verifier) + user_api_key_cache.delete_cache(key=cache_key) + return token_params + + + @staticmethod + def generate_pkce_params() -> Tuple[str, str]: + """ + Generate PKCE (Proof Key for Code Exchange) parameters for OAuth 2.0. + + Returns: + Tuple[str, str]: (code_verifier, code_challenge) + - code_verifier: Random 43-128 character string (we use 43 for efficiency) + - code_challenge: Base64-URL-encoded SHA256 hash of the code_verifier + + Reference: https://datatracker.ietf.org/doc/html/rfc7636 + """ + # Generate a cryptographically random code_verifier (43 characters) + # Using 32 random bytes which becomes 43 characters when base64-url-encoded + code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode('utf-8').rstrip('=') + + # Generate code_challenge using S256 method (SHA256) + code_challenge_bytes = hashlib.sha256(code_verifier.encode('utf-8')).digest() + code_challenge = base64.urlsafe_b64encode(code_challenge_bytes).decode('utf-8').rstrip('=') + + return code_verifier, code_challenge + class MicrosoftSSOHandler: