From aac04d5665374d9369701d892ad156c5f1987d87 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Tue, 3 Feb 2026 07:37:19 +0530 Subject: [PATCH 01/50] fix: SSO PKCE support fails in multi-pod Kubernetes deployments --- litellm/proxy/management_endpoints/ui_sso.py | 92 ++-- .../proxy/management_endpoints/test_ui_sso.py | 467 ++++++++++++------ 2 files changed, 384 insertions(+), 175 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 2d248dc81f3..a2d7510c344 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -11,10 +11,11 @@ Has all /sso/* routes import asyncio import base64 import hashlib +import json import os import secrets from copy import deepcopy -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import RedirectResponse @@ -82,7 +83,15 @@ from litellm.proxy.utils import ( get_server_root_path, ) from litellm.secret_managers.main import get_secret_bool, str_to_bool -from litellm.types.proxy.management_endpoints.ui_sso import * +from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + MicrosoftGraphAPIUserGroupDirectoryObject, + MicrosoftGraphAPIUserGroupResponse, + MicrosoftServicePrincipalTeam, + RoleMappings, + TeamMappings, +) +from litellm.types.proxy.management_endpoints.ui_sso import * # noqa: F403, F401 from litellm.types.proxy.ui_sso import ParsedOpenIDResult if TYPE_CHECKING: @@ -96,15 +105,15 @@ router = APIRouter() def normalize_email(email: Optional[str]) -> Optional[str]: """ Normalize email address to lowercase for consistent storage and comparison. - + Email addresses should be treated as case-insensitive for SSO purposes, even though RFC 5321 technically allows case-sensitive local parts. This prevents issues where SSO providers return emails with different casing than what's stored in the database. - + Args: email: Email address to normalize, can be None - + Returns: Lowercased email address, or None if input is None """ @@ -280,7 +289,7 @@ async def google_login( # check if user defined a custom auth sso sign in handler, if yes, use it if user_custom_ui_sso_sign_in_handler is not None: try: - from litellm_enterprise.proxy.auth.custom_sso_handler import ( + from litellm_enterprise.proxy.auth.custom_sso_handler import ( # type: ignore[import-untyped] EnterpriseCustomSSOHandler, ) @@ -428,7 +437,9 @@ def generic_response_convertor( display_name=get_nested_value( response, generic_user_display_name_attribute_name ), - email=normalize_email(get_nested_value(response, generic_user_email_attribute_name)), + email=normalize_email( + get_nested_value(response, generic_user_email_attribute_name) + ), first_name=get_nested_value(response, generic_user_first_name_attribute_name), last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), @@ -517,6 +528,7 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: if team_mappings_data: from litellm.types.proxy.management_endpoints.ui_sso import TeamMappings + if isinstance(team_mappings_data, dict): team_mappings = TeamMappings(**team_mappings_data) elif isinstance(team_mappings_data, TeamMappings): @@ -554,6 +566,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: if role_mappings_data: from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): role_mappings = RoleMappings(**role_mappings_data) elif isinstance(role_mappings_data, RoleMappings): @@ -567,7 +580,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: verbose_proxy_logger.debug( f"Could not load role_mappings from database: {e}. Continuing with existing role logic." ) - + generic_role_mappings = os.getenv("GENERIC_ROLE_MAPPINGS_ROLES", None) generic_role_mappings_group_claim = os.getenv( "GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None @@ -577,8 +590,8 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: ) if generic_role_mappings is not None: verbose_proxy_logger.debug( - "Found role_mappings for generic provider in environment variables" - ) + "Found role_mappings for generic provider in environment variables" + ) import ast try: @@ -603,7 +616,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: ) return role_mappings except TypeError as e: - verbose_proxy_logger.warning(f"Error decoding role mappings from environment variables: {e}. Continuing with existing role logic.") + verbose_proxy_logger.warning( + f"Error decoding role mappings from environment variables: {e}. Continuing with existing role logic." + ) return role_mappings @@ -875,7 +890,7 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid - """ + """ update_data: dict = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid @@ -1673,7 +1688,7 @@ class SSOAuthenticationHandler: """ from urllib.parse import parse_qs, urlencode, urlparse, urlunparse - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache with generic_sso: # TODO: state should be a random string and added to the user session with cookie @@ -1702,13 +1717,21 @@ class SSOAuthenticationHandler: # If PKCE is enabled, add PKCE parameters to the redirect URL if code_verifier and "state" in redirect_params: - # Store code_verifier in cache (10 min TTL) + # Store code_verifier in cache (10 min TTL). Use Redis when available + # so callbacks landing on another pod can retrieve it (multi-pod SSO). cache_key = f"pkce_verifier:{redirect_params['state']}" - user_api_key_cache.set_cache( - key=cache_key, - value=code_verifier, - ttl=600, - ) + if redis_usage_cache is not None: + redis_usage_cache.set_cache( + key=cache_key, + value=json.dumps(code_verifier), + ttl=600, + ) + else: + user_api_key_cache.set_cache( + key=cache_key, + value=code_verifier, + ttl=600, + ) # Add PKCE parameters to the authorization URL if pkce_params: @@ -2319,27 +2342,38 @@ class SSOAuthenticationHandler: Returns: dict: Token exchange parameters """ - # Prepare token exchange parameters - token_params = {"include_client_id": generic_include_client_id} + # Prepare token exchange parameters (may add code_verifier: str later) + token_params: Dict[str, Any] = {"include_client_id": generic_include_client_id} - # Retrieve PKCE code_verifier if PKCE was used in authorization + # Retrieve PKCE code_verifier if PKCE was used in authorization. + # Use same cache as store: Redis when available (multi-pod), else in-memory. query_params = dict(request.query_params) state = query_params.get("state") if state: - from litellm.proxy.proxy_server import user_api_key_cache + from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache cache_key = f"pkce_verifier:{state}" - code_verifier = user_api_key_cache.get_cache(key=cache_key) + if redis_usage_cache is not None: + code_verifier = redis_usage_cache.get_cache(key=cache_key) + else: + 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 + # Add code_verifier to token exchange parameters (Redis returns decoded string) + token_params["code_verifier"] = ( + code_verifier + if isinstance(code_verifier, str) + else str(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) + if redis_usage_cache is not None: + redis_usage_cache.delete_cache(key=cache_key) + else: + user_api_key_cache.delete_cache(key=cache_key) return token_params @staticmethod @@ -2482,7 +2516,9 @@ class MicrosoftSSOHandler: response = response or {} verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}") openid_response = CustomOpenID( - email=normalize_email(response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail")), + email=normalize_email( + response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail") + ), display_name=response.get(MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE), provider="microsoft", id=response.get(MICROSOFT_USER_ID_ATTRIBUTE), 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 41096503a2e..3a5505cdbc5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2,12 +2,10 @@ import asyncio import json import os import sys -from typing import Optional, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import Request -from fastapi.testclient import TestClient from litellm._uuid import uuid @@ -16,7 +14,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from litellm.proxy._types import LiteLLM_UserTable, NewTeamRequest, NewUserResponse +from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -134,16 +132,32 @@ def test_microsoft_sso_handler_openid_from_response_with_custom_attributes(): expected_team_ids = ["team1"] # Act - with patch("litellm.constants.MICROSOFT_USER_EMAIL_ATTRIBUTE", "custom_email_field"), \ - patch("litellm.constants.MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "custom_display_name"), \ - patch("litellm.constants.MICROSOFT_USER_ID_ATTRIBUTE", "custom_id_field"), \ - patch("litellm.constants.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name"), \ - patch("litellm.constants.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name"), \ - patch("litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_EMAIL_ATTRIBUTE", "custom_email_field"), \ - patch("litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "custom_display_name"), \ - patch("litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_ID_ATTRIBUTE", "custom_id_field"), \ - patch("litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name"), \ - patch("litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name"): + with patch( + "litellm.constants.MICROSOFT_USER_EMAIL_ATTRIBUTE", "custom_email_field" + ), patch( + "litellm.constants.MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "custom_display_name" + ), patch( + "litellm.constants.MICROSOFT_USER_ID_ATTRIBUTE", "custom_id_field" + ), patch( + "litellm.constants.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "custom_first_name" + ), patch( + "litellm.constants.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "custom_last_name" + ), patch( + "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_EMAIL_ATTRIBUTE", + "custom_email_field", + ), patch( + "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", + "custom_display_name", + ), patch( + "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_ID_ATTRIBUTE", + "custom_id_field", + ), patch( + "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", + "custom_first_name", + ), patch( + "litellm.proxy.management_endpoints.ui_sso.MICROSOFT_USER_LAST_NAME_ATTRIBUTE", + "custom_last_name", + ): result = MicrosoftSSOHandler.openid_from_response( response=mock_response, team_ids=expected_team_ids, user_role=None ) @@ -229,7 +243,6 @@ def test_get_microsoft_callback_response_raw_sso_response(): ) # Assert - print("result from verify_and_process", result) assert isinstance(result, dict) assert result["mail"] == "microsoft_user@example.com" assert result["displayName"] == "Microsoft User" @@ -453,10 +466,6 @@ async def test_default_team_params(team_params): # Assert # Verify team was created with correct parameters mock_prisma.db.litellm_teamtable.create.assert_called_once() - print( - "mock_prisma.db.litellm_teamtable.create.call_args", - mock_prisma.db.litellm_teamtable.create.call_args, - ) create_call_args = mock_prisma.db.litellm_teamtable.create.call_args.kwargs[ "data" ] @@ -581,7 +590,7 @@ def test_apply_user_info_values_to_sso_user_defined_values_with_models(): def test_apply_user_info_values_sso_role_takes_precedence(): """ Test that SSO role takes precedence over DB role. - + When Microsoft SSO returns a user_role, it should be used instead of the role stored in the database. This ensures SSO is the authoritative source for user roles. """ @@ -676,16 +685,16 @@ def test_normalize_email(): """ # Test with lowercase email assert normalize_email("test@example.com") == "test@example.com" - + # Test with uppercase email assert normalize_email("TEST@EXAMPLE.COM") == "test@example.com" - + # Test with mixed case email assert normalize_email("Test.User@Example.COM") == "test.user@example.com" - + # Test with None assert normalize_email(None) is None - + # Test with empty string assert normalize_email("") == "" @@ -898,7 +907,7 @@ async def test_upsert_sso_user_no_role_in_sso_response(): def test_get_user_email_and_id_extracts_microsoft_role(): """ Test that _get_user_email_and_id_from_result extracts user_role from Microsoft SSO. - + This ensures Microsoft SSO roles (from app_roles in id_token) are properly extracted and converted from enum to string. """ @@ -964,7 +973,7 @@ async def test_get_user_info_from_db_user_exists(): with patch( "litellm.proxy.management_endpoints.ui_sso.get_user_object" ) as mock_get_user_object: - user_info = await get_user_info_from_db(**args) + await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() assert mock_get_user_object.call_args.kwargs["user_id"] == "krrishd" @@ -1006,7 +1015,7 @@ async def test_get_user_info_from_db_user_exists_alternate_user_id(): with patch( "litellm.proxy.management_endpoints.ui_sso.get_user_object" ) as mock_get_user_object: - user_info = await get_user_info_from_db(**args) + await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() assert mock_get_user_object.call_args.kwargs["user_id"] == "krrishd-email1234" @@ -1015,7 +1024,7 @@ async def test_get_user_info_from_db_user_exists_alternate_user_id(): async def test_get_user_info_from_db_user_not_exists_creates_user(): """ Test that get_user_info_from_db creates a new user when user doesn't exist in DB. - + When get_existing_user_info_from_db returns None, get_user_info_from_db should: 1. Call upsert_sso_user with user_info=None 2. upsert_sso_user should call insert_sso_user to create the user @@ -1103,7 +1112,7 @@ async def test_get_user_info_from_db_user_not_exists_creates_user(): async def test_get_user_info_from_db_user_exists_updates_user(): """ Test that get_user_info_from_db updates existing user when user exists in DB. - + When get_existing_user_info_from_db returns a user, get_user_info_from_db should: 1. Call upsert_sso_user with the existing user_info 2. upsert_sso_user should update the user in the database @@ -1195,6 +1204,7 @@ async def test_get_user_info_from_db_user_exists_updates_user(): # Should return the updated user assert user_info == updated_user + @pytest.mark.asyncio async def test_check_and_update_if_proxy_admin_id(): """ @@ -1302,10 +1312,10 @@ async def test_get_generic_sso_response_with_additional_headers(): mock_sso_class = MagicMock(return_value=mock_sso_instance) with patch.dict(os.environ, test_env_vars): - with patch("fastapi_sso.sso.base.DiscoveryDocument") as mock_discovery: + with patch("fastapi_sso.sso.base.DiscoveryDocument"): with patch( "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class - ) as mock_create_provider: + ): # Act result, received_response = await get_generic_sso_response( request=mock_request, @@ -1363,10 +1373,10 @@ async def test_get_generic_sso_response_with_empty_headers(): mock_sso_class = MagicMock(return_value=mock_sso_instance) with patch.dict(os.environ, test_env_vars): - with patch("fastapi_sso.sso.base.DiscoveryDocument") as mock_discovery: + with patch("fastapi_sso.sso.base.DiscoveryDocument"): with patch( "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class - ) as mock_create_provider: + ): # Act result, received_response = await get_generic_sso_response( request=mock_request, @@ -1751,8 +1761,6 @@ class TestCustomUISSO: """Test that proper error is raised when enterprise module is not available""" from unittest.mock import MagicMock, patch - from litellm.proxy.management_endpoints.ui_sso import google_login - # Mock request mock_request = MagicMock() mock_request.base_url = "https://test.example.com/" @@ -1774,7 +1782,7 @@ class TestCustomUISSO: # This mimics the relevant part of google_login that would trigger the import error try: from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import ( - EnterpriseCustomSSOHandler, + EnterpriseCustomSSOHandler, # noqa: F401 ) return "success" @@ -1978,59 +1986,56 @@ class TestCLIKeyRegenerationFlow: # Test data session_key = "sk-session-456" - + # Mock user info mock_user_info = LiteLLM_UserTable( user_id="test-user-123", user_role="internal_user", teams=["team1", "team2"], - models=["gpt-4"] + models=["gpt-4"], ) # Mock SSO result - mock_sso_result = { - "user_email": "test@example.com", - "user_id": "test-user-123" - } + mock_sso_result = {"user_email": "test@example.com", "user_id": "test-user-123"} # Mock cache mock_cache = MagicMock() - + with patch( "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", - return_value=mock_user_info - ), patch( - "litellm.proxy.proxy_server.prisma_client", MagicMock() - ), patch( + return_value=mock_user_info, + ), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( "litellm.proxy.proxy_server.user_api_key_cache", mock_cache ), patch( "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success", ): - # Act result = await cli_sso_callback( - request=mock_request, key=session_key, existing_key=None, result=mock_sso_result + request=mock_request, + key=session_key, + existing_key=None, + result=mock_sso_result, ) # Assert - verify session was stored in cache mock_cache.set_cache.assert_called_once() call_args = mock_cache.set_cache.call_args - + # Verify cache key format assert "cli_sso_session:" in call_args.kwargs["key"] assert session_key in call_args.kwargs["key"] - + # Verify session data structure session_data = call_args.kwargs["value"] assert session_data["user_id"] == "test-user-123" assert session_data["user_role"] == "internal_user" assert session_data["teams"] == ["team1", "team2"] assert session_data["models"] == ["gpt-4"] - + # Verify TTL assert call_args.kwargs["ttl"] == 600 # 10 minutes - + assert result.status_code == 200 # Verify response contains success message (response is HTML) assert result.body is not None @@ -2046,17 +2051,14 @@ class TestCLIKeyRegenerationFlow: "user_id": "test-user-456", "user_role": "internal_user", "teams": ["team-a", "team-b", "team-c"], - "models": ["gpt-4"] + "models": ["gpt-4"], } # Mock cache mock_cache = MagicMock() mock_cache.get_cache.return_value = session_data - - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", mock_cache - ): + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): # Act - First poll without team_id result = await cli_poll_key(key_id=session_key, team_id=None) @@ -2066,7 +2068,7 @@ class TestCLIKeyRegenerationFlow: assert result["user_id"] == "test-user-456" assert result["teams"] == ["team-a", "team-b", "team-c"] assert "key" not in result # JWT should not be generated yet - + # Verify session was NOT deleted mock_cache.delete_cache.assert_not_called() @@ -2170,34 +2172,33 @@ class TestCLIKeyRegenerationFlow: "user_role": "internal_user", "teams": ["team-a", "team-b", "team-c"], "models": ["gpt-4"], - "user_email": "test@example.com" + "user_email": "test@example.com", } - + # Mock user info mock_user_info = LiteLLM_UserTable( user_id="test-user-789", user_role="internal_user", teams=["team-a", "team-b", "team-c"], - models=["gpt-4"] + models=["gpt-4"], ) # Mock cache mock_cache = MagicMock() mock_cache.get_cache.return_value = session_data - + mock_jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.token" - - with patch( - "litellm.proxy.proxy_server.user_api_key_cache", mock_cache - ), patch( + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), patch( "litellm.proxy.proxy_server.prisma_client" ) as mock_prisma, patch( "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", - return_value=mock_jwt_token + return_value=mock_jwt_token, ) as mock_get_jwt: - # Mock the user lookup - mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user_info) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=mock_user_info + ) # Act - Second poll with team_id result = await cli_poll_key(key_id=session_key, team_id=selected_team) @@ -2208,12 +2209,12 @@ class TestCLIKeyRegenerationFlow: assert result["user_id"] == "test-user-789" assert result["team_id"] == selected_team assert result["teams"] == ["team-a", "team-b", "team-c"] - + # Verify JWT was generated with correct team mock_get_jwt.assert_called_once() jwt_call_args = mock_get_jwt.call_args assert jwt_call_args.kwargs["team_id"] == selected_team - + # Verify session was deleted after JWT generation mock_cache.delete_cache.assert_called_once() @@ -2223,7 +2224,6 @@ class TestGetAppRolesFromIdToken: def test_roles_picked_when_app_roles_not_exists(self): """Test that 'roles' is picked when 'app_roles' doesn't exist""" - import jwt # Create a token with only 'roles' claim token_payload = { @@ -2247,7 +2247,6 @@ class TestGetAppRolesFromIdToken: def test_app_roles_picked_when_both_exist(self): """Test that 'app_roles' takes precedence when both 'app_roles' and 'roles' exist""" - import jwt # Create a token with both 'app_roles' and 'roles' claims token_payload = { @@ -2268,7 +2267,6 @@ class TestGetAppRolesFromIdToken: def test_roles_picked_when_app_roles_is_empty(self): """Test that 'roles' is picked when 'app_roles' exists but is empty""" - import jwt # Create a token with empty 'app_roles' and populated 'roles' token_payload = { @@ -2289,7 +2287,6 @@ class TestGetAppRolesFromIdToken: def test_empty_list_when_neither_exists(self): """Test that empty list is returned when neither 'app_roles' nor 'roles' exist""" - import jwt # Create a token without roles claims token_payload = {"sub": "user123", "email": "test@example.com"} @@ -2313,7 +2310,6 @@ class TestGetAppRolesFromIdToken: def test_empty_list_when_roles_not_a_list(self): """Test that empty list is returned when roles is not a list""" - import jwt # Create a token with non-list roles token_payload = { @@ -2333,7 +2329,6 @@ class TestGetAppRolesFromIdToken: def test_error_handling_on_jwt_decode_exception(self): """Test that exceptions during JWT decode are handled gracefully""" - import jwt mock_token = "invalid.jwt.token" @@ -2726,12 +2721,6 @@ class TestGenericResponseConvertorNestedAttributes: # to handle dotted paths like "attributes.userId" # Current behavior: returns None for nested paths - print(f"User ID result: {result.id}") - print(f"Email result: {result.email}") - print(f"First name result: {result.first_name}") - print(f"Last name result: {result.last_name}") - print(f"Display name result: {result.display_name}") - # Expected behavior with current implementation (no nested path support): assert result.id == "nested-user-456" assert ( @@ -2821,14 +2810,15 @@ class TestGetGenericSSORedirectParams: # Arrange cli_state = "litellm-session-token:sk-test123" - + with patch.dict(os.environ, {"GENERIC_CLIENT_STATE": "env_state_value"}): # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=cli_state, - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=cli_state, + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert @@ -2843,14 +2833,15 @@ class TestGetGenericSSORedirectParams: # Arrange env_state = "custom_env_state_value" - + with patch.dict(os.environ, {"GENERIC_CLIENT_STATE": env_state}): # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=None, - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=None, + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert @@ -2867,13 +2858,14 @@ class TestGetGenericSSORedirectParams: with patch.dict(os.environ, {}, clear=False): # Remove GENERIC_CLIENT_STATE if it exists os.environ.pop("GENERIC_CLIENT_STATE", None) - + # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=None, - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=None, + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert @@ -2893,26 +2885,27 @@ class TestGetGenericSSORedirectParams: # Arrange test_state = "test_state_123" - + with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=test_state, - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=test_state, + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert state assert redirect_params["state"] == test_state - + # Assert PKCE parameters assert code_verifier is not None assert len(code_verifier) == 43 # Standard PKCE verifier length assert "code_challenge" in redirect_params assert "code_challenge_method" in redirect_params assert redirect_params["code_challenge_method"] == "S256" - + # Verify code_challenge is correctly derived from code_verifier expected_challenge_bytes = hashlib.sha256( code_verifier.encode("utf-8") @@ -2932,14 +2925,15 @@ class TestGetGenericSSORedirectParams: # Arrange test_state = "test_state_456" - + with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "false"}): # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=test_state, - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=test_state, + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert @@ -2957,7 +2951,7 @@ class TestGetGenericSSORedirectParams: # Arrange cli_state = "cli_state_priority" env_state = "env_state_should_not_be_used" - + with patch.dict( os.environ, { @@ -2966,17 +2960,18 @@ class TestGetGenericSSORedirectParams: }, ): # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state=cli_state, - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=cli_state, + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert assert redirect_params["state"] == cli_state # CLI state takes priority assert redirect_params["state"] != env_state - + # PKCE should still be generated assert code_verifier is not None assert "code_challenge" in redirect_params @@ -2990,14 +2985,15 @@ class TestGetGenericSSORedirectParams: # Arrange env_state = "env_state_for_empty_cli" - + with patch.dict(os.environ, {"GENERIC_CLIENT_STATE": env_state}): # Act - redirect_params, code_verifier = ( - SSOAuthenticationHandler._get_generic_sso_redirect_params( - state="", # Empty string - generic_authorization_endpoint="https://auth.example.com/authorize", - ) + ( + redirect_params, + code_verifier, + ) = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state="", # Empty string + generic_authorization_endpoint="https://auth.example.com/authorize", ) # Assert - empty string is falsy, so env variable should be used @@ -3014,7 +3010,7 @@ class TestGetGenericSSORedirectParams: # Arrange - no state provided with patch.dict(os.environ, {}, clear=False): os.environ.pop("GENERIC_CLIENT_STATE", None) - + # Act params1, _ = SSOAuthenticationHandler._get_generic_sso_redirect_params( state=None, @@ -3145,6 +3141,176 @@ class TestPKCEFunctionality: assert "code_challenge_method=S256" in updated_location assert f"state={test_state}" in updated_location + @pytest.mark.asyncio + async def test_pkce_redis_multi_pod_verifier_roundtrip(self): + """ + Mock Redis to verify PKCE code_verifier round-trip across "pods": + Pod A stores verifier in Redis; Pod B retrieves it (no real IdP). + """ + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # In-memory mock of Redis (shared between "pods") + class MockRedisCache: + def __init__(self): + self._store = {} + + def set_cache(self, key, value, **kwargs): + self._store[key] = value + + def get_cache(self, key, **kwargs): + val = self._store.get(key) + if val is None: + return None + # Simulate RedisCache._get_cache_logic: stored as JSON string, return decoded + if isinstance(val, str): + try: + return json.loads(val) + except (ValueError, TypeError): + return val + return val + + def delete_cache(self, key): + self._store.pop(key, None) + + mock_redis = MockRedisCache() + mock_in_memory = MagicMock() + + mock_sso = MagicMock() + mock_redirect_response = MagicMock() + mock_redirect_response.headers = { + "location": "https://auth.example.com/authorize?state=multi_pod_state_xyz&client_id=abc" + } + mock_sso.get_login_redirect = AsyncMock(return_value=mock_redirect_response) + mock_sso.__enter__ = MagicMock(return_value=mock_sso) + mock_sso.__exit__ = MagicMock(return_value=False) + + with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): + with patch("litellm.proxy.proxy_server.redis_usage_cache", mock_redis): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory + ): + # Pod A: start login, store code_verifier in "Redis" + await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_sso, + state="multi_pod_state_xyz", + generic_authorization_endpoint="https://auth.example.com/authorize", + ) + mock_in_memory.set_cache.assert_not_called() + # MockRedisCache is a real class; assert on state, not .assert_called_* + stored_key = "pkce_verifier:multi_pod_state_xyz" + assert stored_key in mock_redis._store + stored_value = mock_redis._store[stored_key] + assert json.loads(stored_value) # valid verifier string + + # Pod B: callback with same state, retrieve from "Redis" + mock_request = MagicMock(spec=Request) + mock_request.query_params = {"state": "multi_pod_state_xyz"} + token_params = ( + SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) + ) + assert "code_verifier" in token_params + assert token_params["code_verifier"] == json.loads(stored_value) + mock_in_memory.get_cache.assert_not_called() + # delete_cache called; key removed (asserted below) + + # Verifier consumed (single-use); key removed from "Redis" + assert "pkce_verifier:multi_pod_state_xyz" not in mock_redis._store + + @pytest.mark.asyncio + async def test_pkce_fallback_in_memory_roundtrip_when_redis_none(self): + """ + Regression: When redis_usage_cache is None (no Redis configured), + code_verifier is stored and retrieved via user_api_key_cache. + Roundtrip works when callback hits same pod (same in-memory cache). + Single-pod or no-Redis deployments must continue to work. + """ + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # In-memory store (simulates user_api_key_cache on one pod) + in_memory_store = {} + + def set_cache(key, value, **kwargs): + in_memory_store[key] = value + + def get_cache(key, **kwargs): + return in_memory_store.get(key) + + def delete_cache(key): + in_memory_store.pop(key, None) + + mock_in_memory = MagicMock() + mock_in_memory.set_cache.side_effect = set_cache + mock_in_memory.get_cache.side_effect = get_cache + mock_in_memory.delete_cache.side_effect = delete_cache + + mock_sso = MagicMock() + mock_redirect_response = MagicMock() + mock_redirect_response.headers = { + "location": "https://auth.example.com/authorize?state=fallback_state_xyz&client_id=abc" + } + mock_sso.get_login_redirect = AsyncMock(return_value=mock_redirect_response) + mock_sso.__enter__ = MagicMock(return_value=mock_sso) + mock_sso.__exit__ = MagicMock(return_value=False) + + with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): + with patch("litellm.proxy.proxy_server.redis_usage_cache", None): + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory + ): + # Pod A: start login, store code_verifier in in-memory cache + await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_sso, + state="fallback_state_xyz", + generic_authorization_endpoint="https://auth.example.com/authorize", + ) + mock_in_memory.set_cache.assert_called_once() + stored_key = mock_in_memory.set_cache.call_args.kwargs["key"] + stored_value = mock_in_memory.set_cache.call_args.kwargs["value"] + assert stored_key == "pkce_verifier:fallback_state_xyz" + assert isinstance(stored_value, str) and len(stored_value) == 43 + + # Same pod: callback retrieves from in-memory cache + mock_request = MagicMock(spec=Request) + mock_request.query_params = {"state": "fallback_state_xyz"} + token_params = ( + SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) + ) + assert "code_verifier" in token_params + assert token_params["code_verifier"] == stored_value + mock_in_memory.get_cache.assert_called_once_with(key=stored_key) + mock_in_memory.delete_cache.assert_called_once_with(key=stored_key) + + # Verifier consumed; key removed from in-memory + assert "pkce_verifier:fallback_state_xyz" not in in_memory_store + + @pytest.mark.asyncio + async def test_pkce_prepare_token_exchange_returns_nothing_when_no_state(self): + """ + Regression: prepare_token_exchange_parameters with no state in request + does not call cache and does not add code_verifier. + """ + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + mock_redis = MagicMock() + mock_in_memory = MagicMock() + + with patch("litellm.proxy.proxy_server.redis_usage_cache", mock_redis): + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_in_memory): + mock_request = MagicMock(spec=Request) + mock_request.query_params = {} + token_params = ( + SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) + ) + assert "code_verifier" not in token_params + mock_redis.get_cache.assert_not_called() + mock_in_memory.get_cache.assert_not_called() + # Tests for SSO user team assignment bug (Issue: SSO Users Not Added to Entra-Synced Teams on First Login) class TestAddMissingTeamMember: @@ -3268,9 +3434,7 @@ class TestAddMissingTeamMember: team_member_calls = [] async def track_team_member_add(team_id, user_info): - team_member_calls.append( - {"team_id": team_id, "user_id": user_info.user_id} - ) + team_member_calls.append({"team_id": team_id, "user_id": user_info.user_id}) # New SSO user with Entra groups new_user = NewUserResponse( @@ -3331,7 +3495,6 @@ class TestAddMissingTeamMember: """ Parametrized test ensuring add_missing_team_member works for all user types. """ - from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member user_info = user_info_factory("test-user-id") @@ -3421,7 +3584,7 @@ async def test_role_mappings_override_default_internal_user_params(): return_value=mock_new_user_response, ) as mock_new_user: # Act - result = await insert_sso_user( + _ = await insert_sso_user( result_openid=mock_result_openid, user_defined_values=user_defined_values, ) @@ -3443,7 +3606,7 @@ async def test_role_mappings_override_default_internal_user_params(): assert ( new_user_request.budget_duration == "30d" ), "budget_duration from default_internal_user_params should be applied" - + # Note: models are applied via _update_internal_new_user_params inside new_user, # not in insert_sso_user, so we verify user_defined_values was updated correctly # by checking that the function completed successfully and other defaults were applied @@ -3558,7 +3721,10 @@ class TestSSOReadinessEndpoint: assert data["sso_configured"] is True assert data["provider"] == "google" assert "GOOGLE_CLIENT_SECRET" in data["missing_environment_variables"] - assert "Google SSO is configured but missing required environment variables" in data["message"] + assert ( + "Google SSO is configured but missing required environment variables" + in data["message"] + ) finally: app.dependency_overrides.clear() @@ -3607,7 +3773,7 @@ class TestSSOReadinessEndpoint: response = client.get("/sso/readiness") assert response.status_code == expected_status - + if expected_status == 200: data = response.json() assert data["sso_configured"] is True @@ -3677,7 +3843,7 @@ class TestSSOReadinessEndpoint: response = client.get("/sso/readiness") assert response.status_code == expected_status - + if expected_status == 200: data = response.json() assert data["sso_configured"] is True @@ -3722,8 +3888,14 @@ class TestCustomMicrosoftSSO: discovery = await sso.get_discovery_document() - assert discovery["authorization_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/authorize" - assert discovery["token_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" + assert ( + discovery["authorization_endpoint"] + == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/authorize" + ) + assert ( + discovery["token_endpoint"] + == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" + ) assert discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" @pytest.mark.asyncio @@ -3787,8 +3959,13 @@ class TestCustomMicrosoftSSO: # Custom auth endpoint assert discovery["authorization_endpoint"] == custom_auth_endpoint # Default token and userinfo endpoints - assert discovery["token_endpoint"] == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" - assert discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" + assert ( + discovery["token_endpoint"] + == "https://login.microsoftonline.com/test-tenant/oauth2/v2.0/token" + ) + assert ( + discovery["userinfo_endpoint"] == "https://graph.microsoft.com/v1.0/me" + ) def test_custom_microsoft_sso_uses_common_tenant_when_none(self): """ @@ -3825,11 +4002,7 @@ async def test_setup_team_mappings(): # Arrange mock_prisma = MagicMock() mock_sso_config = MagicMock() - mock_sso_config.sso_settings = { - "team_mappings": { - "team_ids_jwt_field": "groups" - } - } + mock_sso_config.sso_settings = {"team_mappings": {"team_ids_jwt_field": "groups"}} mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( return_value=mock_sso_config ) From 768f9a44b21a8af72cbeb44d9099697e540aa580 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Tue, 3 Feb 2026 09:50:10 +0530 Subject: [PATCH 02/50] fix: virutal key grace period from env/UI --- docs/my-website/docs/proxy/config_settings.md | 1 + docs/my-website/docs/proxy/virtual_keys.md | 7 +- .../migration.sql | 19 ++ .../litellm_proxy_extras/schema.prisma | 13 + litellm/constants.py | 51 +++- litellm/proxy/_types.py | 4 +- .../common_utils/key_rotation_manager.py | 29 +- .../key_management_endpoints.py | 253 +++++++++++------- litellm/proxy/schema.prisma | 13 + litellm/proxy/utils.py | 65 +++++ schema.prisma | 13 + .../common_utils/test_key_rotation_manager.py | 193 +++++++++---- .../organisms/regenerate_key_modal.tsx | 12 + 13 files changed, 500 insertions(+), 173 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 264c7d765b3..1a1bf975d2d 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -755,6 +755,7 @@ router_settings: | LITELM_ENVIRONMENT | Environment of LiteLLM Instance, used by logging services. Currently only used by DeepEval. | LITELLM_KEY_ROTATION_ENABLED | Enable auto-key rotation for LiteLLM (boolean). Default is false. | LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS | Interval in seconds for how often to run job that auto-rotates keys. Default is 86400 (24 hours). +| LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS | Hours to keep old key valid after rotation (e.g. 24, 48, 72). Default is 0 (immediate revoke). Used for scheduled rotations and as fallback when not specified in regenerate request. | LITELLM_LICENSE | License key for LiteLLM usage | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM | LITELLM_LOG | Enable detailed logging for LiteLLM diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 38ff4ede280..31506ea2611 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -549,11 +549,14 @@ curl 'http://localhost:4000/key/sk-1234/regenerate' \ "models": [ "gpt-4", "gpt-3.5-turbo" - ] + ], + "grace_period_hours": 48 }' ``` +**Grace period (optional)**: Set `grace_period_hours` (e.g. 24, 48, 72) to keep the old key valid for a transitional period. Both old and new keys work until the grace period elapses, enabling seamless cutover without production downtime. Default is 0 (immediate revoke). Can also be set via `LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS` for scheduled rotations. + **Read More** - [Write rotated keys to secrets manager](https://docs.litellm.ai/docs/secret#aws-secret-manager) @@ -640,11 +643,13 @@ Set these environment variables when starting the proxy: |----------|-------------|---------| | `LITELLM_KEY_ROTATION_ENABLED` | Enable the rotation worker | `false` | | `LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS` | How often to scan for keys to rotate (in seconds) | `86400` (24 hours) | +| `LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS` | Hours to keep old key valid after rotation (24–72 recommended) | `0` (immediate revoke) | **Example:** ```bash export LITELLM_KEY_ROTATION_ENABLED=true export LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS=3600 # Check every hour +export LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS=48 # Keep old key valid for 48h during cutover litellm --config config.yaml ``` diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql new file mode 100644 index 00000000000..51d88444191 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260203120000_add_deprecated_verification_token_table/migration.sql @@ -0,0 +1,19 @@ +-- CreateTable +CREATE TABLE "LiteLLM_DeprecatedVerificationToken" ( + "id" TEXT NOT NULL, + "token" TEXT NOT NULL, + "active_token_id" TEXT NOT NULL, + "revoke_at" TIMESTAMP(3) NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_DeprecatedVerificationToken_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_DeprecatedVerificationToken_token_key" ON "LiteLLM_DeprecatedVerificationToken"("token"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DeprecatedVerificationToken_token_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("token", "revoke_at"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DeprecatedVerificationToken_revoke_at_idx" ON "LiteLLM_DeprecatedVerificationToken"("revoke_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 3b81da10923..8992f4d8df6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -317,6 +317,19 @@ model LiteLLM_VerificationToken { @@index([budget_reset_at, expires]) } +// Deprecated keys during grace period - allows old key to work until revoke_at +model LiteLLM_DeprecatedVerificationToken { + id String @id @default(uuid()) + token String // Hashed old key + active_token_id String // Current token hash in LiteLLM_VerificationToken + revoke_at DateTime // When the old key stops working + created_at DateTime @default(now()) @map("created_at") + + @@unique([token]) + @@index([token, revoke_at]) + @@index([revoke_at]) +} + // Audit table for deleted keys - preserves spend and key information for historical tracking model LiteLLM_DeletedVerificationToken { id String @id @default(uuid()) diff --git a/litellm/constants.py b/litellm/constants.py index 6427c367924..7625d55f003 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -125,15 +125,19 @@ _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client fo # Aiohttp connection pooling - prevents memory leaks from unbounded connection growth # Set to 0 for unlimited (not recommended for production) AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300)) -AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50)) +AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int( + os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50) +) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) # enable_cleanup_closed is only needed for Python versions with the SSL leak bug # Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960) # Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78 -AIOHTTP_NEEDS_CLEANUP_CLOSED = ( - (3, 13, 0) <= sys.version_info < (3, 13, 1) or sys.version_info < (3, 12, 7) -) +AIOHTTP_NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < ( + 3, + 13, + 1, +) or sys.version_info < (3, 12, 7) # WebSocket constants # Default to None (unlimited) to match OpenAI's official agents SDK behavior @@ -171,7 +175,9 @@ REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer" -REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer" +REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = ( + "litellm_daily_end_user_spend_update_buffer" +) REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) @@ -297,7 +303,9 @@ MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int( DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000)) #### Networking settings #### request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds -DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes +DEFAULT_A2A_AGENT_TIMEOUT: float = float( + os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000) +) # 10 minutes STREAM_SSE_DONE_STRING: str = "[DONE]" STREAM_SSE_DATA_PREFIX: str = "data: " ### SPEND TRACKING ### @@ -333,8 +341,12 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) -EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds -EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)) # 80% of max budget +EMAIL_BUDGET_ALERT_TTL = int( + os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60) +) # 24 hours in seconds +EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float( + os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8) +) # 80% of max budget ############### LLM Provider Constants ############### ### ANTHROPIC CONSTANTS ### ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv( @@ -1082,7 +1094,17 @@ known_tokenizer_config = { } -OPENAI_FINISH_REASONS = ["stop", "length", "function_call", "content_filter", "null", "finish_reason_unspecified", "malformed_function_call", "guardrail_intervened", "eos"] +OPENAI_FINISH_REASONS = [ + "stop", + "length", + "function_call", + "content_filter", + "null", + "finish_reason_unspecified", + "malformed_function_call", + "guardrail_intervened", + "eos", +] HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int( os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60) ) # 1 minute @@ -1172,6 +1194,9 @@ LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false" LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int( os.getenv("LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400) ) # 24 hours default +LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS = int( + os.getenv("LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS", 0) +) # Hours to keep old key valid after rotation; 0 = immediate revoke (default) UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" @@ -1182,8 +1207,8 @@ CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session" CLI_JWT_TOKEN_NAME = "cli-jwt-token" # Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility CLI_JWT_EXPIRATION_HOURS = int( - os.getenv("CLI_JWT_EXPIRATION_HOURS") - or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") + os.getenv("CLI_JWT_EXPIRATION_HOURS") + or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") or 24 ) @@ -1361,9 +1386,7 @@ MICROSOFT_USER_EMAIL_ATTRIBUTE = str( MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str( os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName") ) -MICROSOFT_USER_ID_ATTRIBUTE = str( - os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id") -) +MICROSOFT_USER_ID_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id")) MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str( os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName") ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9ae95085f55..dd3f83713f8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -359,7 +359,6 @@ class LiteLLMRoutes(enum.Enum): "/v1/vector_stores/{vector_store_id}/files/{file_id}/content", "/vector_store/list", "/v1/vector_store/list", - # search "/search", "/v1/search", @@ -986,6 +985,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None + grace_period_hours: Optional[ + int + ] = None # Hours to keep old key valid; 0/None = immediate revoke class ResetSpendRequest(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 13bbf2272f7..7ef26bf07b9 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -8,7 +8,10 @@ from datetime import datetime, timezone from typing import List from litellm._logging import verbose_proxy_logger -from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME +from litellm.constants import ( + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS, +) from litellm.proxy._types import ( GenerateKeyResponse, LiteLLM_VerificationToken, @@ -37,6 +40,9 @@ class KeyRotationManager: try: verbose_proxy_logger.info("Starting scheduled key rotation check...") + # Clean up expired deprecated keys first + await self._cleanup_expired_deprecated_keys() + # Find keys that are due for rotation keys_to_rotate = await self._find_keys_needing_rotation() @@ -97,6 +103,24 @@ class KeyRotationManager: return keys_with_rotation + async def _cleanup_expired_deprecated_keys(self) -> None: + """ + Remove deprecated key entries whose revoke_at has passed. + """ + try: + now = datetime.now(timezone.utc) + result = await self.prisma_client.db.litellm_deprecatedverificationtoken.delete_many( + where={"revoke_at": {"lt": now}} + ) + if result.get("count", 0) > 0: + verbose_proxy_logger.debug( + "Cleaned up %s expired deprecated key(s)", result["count"] + ) + except Exception as e: + verbose_proxy_logger.debug( + "Deprecated key cleanup skipped (table may not exist): %s", e + ) + def _should_rotate_key(self, key: LiteLLM_VerificationToken, now: datetime) -> bool: """ Determine if a key should be rotated based on key_rotation_at timestamp. @@ -115,10 +139,11 @@ class KeyRotationManager: """ Rotate a single key using existing regenerate_key_fn and call the rotation hook """ - # Create regenerate request + # Create regenerate request with grace period for seamless cutover regenerate_request = RegenerateKeyRequest( key=key.token or "", key_alias=key.key_alias, # Pass key alias to ensure correct secret is updated in AWS Secrets Manager + grace_period_hours=LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS, ) # Create a system user for key rotation diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d1840363009..dea1e55f7f9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -12,6 +12,7 @@ All /key management endpoints import asyncio import copy import json +import os import secrets import traceback from datetime import datetime, timedelta, timezone @@ -518,7 +519,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 ) # Handle special case where duration is "-1" (never expires) if value == "-1": - user_duration = float('inf') # Infinite duration + user_duration = float("inf") # Infinite duration else: user_duration = duration_in_seconds(duration=value) if user_duration > upperbound_duration: @@ -660,9 +661,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 request_type="key", **data_json, table_name="key" ) - response["soft_budget"] = ( - data.soft_budget - ) # include the user-input soft budget in the response + response[ + "soft_budget" + ] = data.soft_budget # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1083,12 +1084,16 @@ async def generate_key_fn( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) if data.soft_budget is not None and data.soft_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + detail={ + "error": f"soft_budget cannot be negative. Received: {data.soft_budget}" + }, ) if user_custom_key_generate is not None: @@ -1336,6 +1341,7 @@ async def prepare_key_update_data( data_json: dict = data.model_dump(exclude_unset=True) data_json.pop("key", None) data_json.pop("new_key", None) + data_json.pop("grace_period_hours", None) # Request-only param, not a DB column if ( data.metadata is not None and data.metadata.get("service_account_id") is not None @@ -1399,8 +1405,13 @@ async def prepare_key_update_data( validate_model_max_budget(non_default_values["model_max_budget"]) # Serialize router_settings to JSON if present - if "router_settings" in non_default_values and non_default_values["router_settings"] is not None: - non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"]) + if ( + "router_settings" in non_default_values + and non_default_values["router_settings"] is not None + ): + non_default_values["router_settings"] = safe_dumps( + non_default_values["router_settings"] + ) non_default_values = prepare_metadata_fields( data=data, non_default_values=non_default_values, existing_metadata=_metadata @@ -1448,19 +1459,17 @@ def is_different_team( def _validate_max_budget(max_budget: Optional[float]) -> None: """ Validate that max_budget is not negative. - + Args: max_budget: The max_budget value to validate - + Raises: HTTPException: If max_budget is negative """ if max_budget is not None and max_budget < 0: raise HTTPException( status_code=400, - detail={ - "error": f"max_budget cannot be negative. Received: {max_budget}" - }, + detail={"error": f"max_budget cannot be negative. Received: {max_budget}"}, ) @@ -1469,14 +1478,14 @@ async def _get_and_validate_existing_key( ) -> LiteLLM_VerificationToken: """ Get existing key from database and validate it exists. - + Args: token: The key token to look up prisma_client: Prisma client instance - + Returns: LiteLLM_VerificationToken: The existing key row - + Raises: HTTPException: If key is not found """ @@ -1485,19 +1494,19 @@ async def _get_and_validate_existing_key( status_code=500, detail={"error": "Database not connected"}, ) - + existing_key_row = await prisma_client.get_data( token=token, table_name="key", query_type="find_unique", ) - + if existing_key_row is None: raise HTTPException( status_code=404, detail={"error": f"Key not found: {token}"}, ) - + return existing_key_row @@ -1512,10 +1521,10 @@ async def _process_single_key_update( ) -> Dict[str, Any]: """ Process a single key update with all validations and checks. - + This function encapsulates all the logic for updating a single key, including validation, permission checks, team checks, and database updates. - + Args: key_update_item: The key update request item user_api_key_dict: The authenticated user's API key info @@ -1524,22 +1533,22 @@ async def _process_single_key_update( user_api_key_cache: User API key cache proxy_logging_obj: Proxy logging object llm_router: LLM router instance - + Returns: Dict containing the updated key information - + Raises: HTTPException: For various validation and permission errors """ # Validate max_budget _validate_max_budget(key_update_item.max_budget) - + # Get and validate existing key existing_key_row = await _get_and_validate_existing_key( token=key_update_item.key, prisma_client=prisma_client, ) - + # Check team member permissions if prisma_client is not None: await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( @@ -1549,7 +1558,7 @@ async def _process_single_key_update( existing_key_row=existing_key_row, user_api_key_cache=user_api_key_cache, ) - + # Create UpdateKeyRequest from BulkUpdateKeyRequestItem update_key_request = UpdateKeyRequest( key=key_update_item.key, @@ -1558,7 +1567,7 @@ async def _process_single_key_update( team_id=key_update_item.team_id, tags=key_update_item.tags, ) - + # Get team object and check team limits if team_id is provided team_obj: Optional[LiteLLM_TeamTableCachedObj] = None if update_key_request.team_id is not None: @@ -1568,18 +1577,16 @@ async def _process_single_key_update( user_api_key_cache=user_api_key_cache, check_db_only=True, ) - + if team_obj is not None and prisma_client is not None: await _check_team_key_limits( team_table=team_obj, data=update_key_request, prisma_client=prisma_client, ) - + # Validate team change if team is being changed - if is_different_team( - data=update_key_request, existing_key_row=existing_key_row - ): + if is_different_team(data=update_key_request, existing_key_row=existing_key_row): if llm_router is None: raise HTTPException( status_code=400, @@ -1590,9 +1597,7 @@ async def _process_single_key_update( if team_obj is None: raise HTTPException( status_code=500, - detail={ - "error": "Team object not found for team change validation" - }, + detail={"error": "Team object not found for team change validation"}, ) validate_key_team_change( key=existing_key_row, @@ -1600,31 +1605,29 @@ async def _process_single_key_update( change_initiated_by=user_api_key_dict, llm_router=llm_router, ) - + # Prepare update data non_default_values = await prepare_key_update_data( data=update_key_request, existing_key_row=existing_key_row ) - + # Update key in database if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Database not connected"}, ) - + _data = {**non_default_values, "token": key_update_item.key} - response = await prisma_client.update_data( - token=key_update_item.key, data=_data - ) - + response = await prisma_client.update_data(token=key_update_item.key, data=_data) + # Delete cache await _delete_cache_key_object( hashed_token=hash_token(key_update_item.key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - + # Trigger async hook asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -1635,19 +1638,19 @@ async def _process_single_key_update( litellm_changed_by=litellm_changed_by, ) ) - + if response is None: raise ValueError("Failed to update key got response = None") - + # Extract and format updated key info updated_key_info = response.get("data", {}) if hasattr(updated_key_info, "model_dump"): updated_key_info = updated_key_info.model_dump() elif hasattr(updated_key_info, "dict"): updated_key_info = updated_key_info.dict() - + updated_key_info.pop("token", None) - + return updated_key_info @@ -1740,7 +1743,9 @@ async def update_key_fn( if data.max_budget is not None and data.max_budget < 0: raise HTTPException( status_code=400, - detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + detail={ + "error": f"max_budget cannot be negative. Received: {data.max_budget}" + }, ) data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) @@ -1959,13 +1964,11 @@ async def bulk_update_keys( proxy_logging_obj, user_api_key_cache, ) - + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, - detail={ - "error": "Only proxy admins can perform bulk key updates" - }, + detail={"error": "Only proxy admins can perform bulk key updates"}, ) if prisma_client is None: @@ -2505,7 +2508,9 @@ async def generate_key_helper_fn( # noqa: PLR0915 aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) - router_settings_json = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) + router_settings_json = ( + safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) + ) # Add model_rpm_limit and model_tpm_limit to metadata if model_rpm_limit is not None: @@ -2672,10 +2677,12 @@ async def generate_key_helper_fn( # noqa: PLR0915 ) key_data["created_at"] = getattr(create_key_response, "created_at", None) key_data["updated_at"] = getattr(create_key_response, "updated_at", None) - + # Deserialize router_settings from JSON string to dict for response router_settings_value = key_data.get("router_settings") - if router_settings_value is not None and isinstance(router_settings_value, str): + if router_settings_value is not None and isinstance( + router_settings_value, str + ): try: key_data["router_settings"] = yaml.safe_load(router_settings_value) except yaml.YAMLError: @@ -2758,27 +2765,27 @@ async def can_modify_verification_token( ) -> bool: """ Check if user has permission to modify (delete/regenerate) a verification token. - + Rules: - Proxy admin can modify any key - For team keys: only team admin or key owner can modify - For personal keys: only key owner can modify - + Args: key_info: The verification token to check user_api_key_cache: Cache for user API keys user_api_key_dict: The user making the request prisma_client: Prisma client for database access - + Returns: True if user can modify the key, False otherwise """ is_team_key = _is_team_key(data=key_info) - + # 1. Proxy admin can modify any key if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return True - + # 2. For team keys: only team admin or key owner can modify if is_team_key and key_info.team_id is not None: # Get team object to check if user is team admin @@ -2788,34 +2795,35 @@ async def can_modify_verification_token( user_api_key_cache=user_api_key_cache, check_db_only=True, ) - + if team_table is None: return False - + # Check if user is team admin if _is_user_team_admin( user_api_key_dict=user_api_key_dict, team_obj=team_table, ): return True - + # Check if the key belongs to the user (they own it) - if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: + if ( + key_info.user_id is not None + and key_info.user_id == user_api_key_dict.user_id + ): return True - + # Not team admin and doesn't own the key return False - + # 3. For personal keys: only key owner can modify if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id: return True - + # Default: deny return False - - async def delete_verification_tokens( tokens: List, user_api_key_cache: DualCache, @@ -2845,10 +2853,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} - ) + _keys_being_deleted: List[ + LiteLLM_VerificationToken + ] = await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} ) if len(_keys_being_deleted) == 0: @@ -2948,11 +2956,24 @@ def _transform_verification_tokens_to_deleted_records( if org_id_value is not None: record["organization_id"] = org_id_value - for json_field in ["aliases", "config", "permissions", "metadata", "model_spend", "model_max_budget", "router_settings"]: + for json_field in [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ]: if json_field in record and record[json_field] is not None: record[json_field] = json.dumps(record[json_field]) - for rel_key in ("litellm_budget_table", "litellm_organization_table", "object_permission", "id"): + for rel_key in ( + "litellm_budget_table", + "litellm_organization_table", + "object_permission", + "id", + ): record.pop(rel_key, None) records.append(record) @@ -2967,9 +2988,7 @@ async def _save_deleted_verification_token_records( """Save deleted verification token records to the database.""" if not records: return - await prisma_client.db.litellm_deletedverificationtoken.create_many( - data=records - ) + await prisma_client.db.litellm_deletedverificationtoken.create_many(data=records) async def _persist_deleted_verification_tokens( @@ -3032,9 +3051,9 @@ async def _rotate_master_key( from litellm.proxy.proxy_server import proxy_config try: - models: Optional[List] = ( - await prisma_client.db.litellm_proxymodeltable.find_many() - ) + models: Optional[ + List + ] = await prisma_client.db.litellm_proxymodeltable.find_many() except Exception: models = None # 2. process model table @@ -3111,7 +3130,9 @@ async def _rotate_master_key( updated_patch=decrypted_cred, new_encryption_key=new_master_key, ) - credential_object_jsonified = jsonify_object(encrypted_cred.model_dump()) + credential_object_jsonified = jsonify_object( + encrypted_cred.model_dump() + ) await prisma_client.db.litellm_credentialstable.update( where={"credential_name": cred.credential_name}, data={ @@ -3156,7 +3177,7 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: dependencies=[Depends(user_api_key_auth)], ) @management_endpoint_wrapper -async def regenerate_key_fn( +async def regenerate_key_fn( # noqa: PLR0915 key: Optional[str] = None, data: Optional[RegenerateKeyRequest] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -3195,6 +3216,7 @@ async def regenerate_key_fn( - permissions: Optional[dict] - Key-specific permissions - guardrails: Optional[List[str]] - List of active guardrails for the key - blocked: Optional[bool] - Whether the key is blocked + - grace_period_hours: Optional[int] - Hours to keep old key valid after rotation (e.g. 24, 48, 72). 0 or omitted = immediate revoke. Env: LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS Returns: @@ -3330,6 +3352,37 @@ async def regenerate_key_fn( update_data.update(non_default_values) update_data = prisma_client.jsonify_object(data=update_data) + + # If grace period > 0, insert deprecated key so old key remains valid + if data is not None and data.grace_period_hours is not None: + grace_period_hours = data.grace_period_hours + else: + grace_period_hours = int( + os.getenv("LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS", "0") + ) + if grace_period_hours > 0: + try: + revoke_at = datetime.now(timezone.utc) + timedelta( + hours=grace_period_hours + ) + await prisma_client.db.litellm_deprecatedverificationtoken.create( + data={ + "token": hashed_api_key, + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + } + ) + verbose_proxy_logger.debug( + "Deprecated key retained for %s hours (revoke_at: %s)", + grace_period_hours, + revoke_at, + ) + except Exception as deprecated_err: + verbose_proxy_logger.warning( + "Failed to insert deprecated key for grace period: %s", + deprecated_err, + ) + # Update the token in the database updated_token = await prisma_client.db.litellm_verificationtoken.update( where={"token": hashed_api_key}, @@ -3423,7 +3476,9 @@ def _validate_reset_spend_value( if reset_to > current_spend: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail={"error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})"}, + detail={ + "error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})" + }, ) max_budget = key_in_db.max_budget @@ -3549,11 +3604,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, - ) + complete_user_info_db_obj: Optional[ + BaseModel + ] = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, ) if complete_user_info_db_obj is None: @@ -3639,10 +3694,10 @@ async def get_admin_team_ids( if complete_user_info is None: return [] # Get all teams that user is an admin of - teams: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} - ) + teams: Optional[ + List[BaseModel] + ] = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} ) if teams is None: return [] @@ -3687,8 +3742,12 @@ async def list_keys( description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), - expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), - status: Optional[str] = Query(None, description="Filter by status (e.g. 'deleted')"), + expand: Optional[List[str]] = Query( + None, description="Expand related objects (e.g. 'user')" + ), + status: Optional[str] = Query( + None, description="Filter by status (e.g. 'deleted')" + ), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. @@ -3780,7 +3839,9 @@ async def list_keys( message=getattr(e, "detail", f"error({str(e)})"), type=ProxyErrorTypes.internal_server_error, param=getattr(e, "param", "None"), - code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR), + code=getattr( + e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR + ), ) elif isinstance(e, ProxyException): raise e diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3b81da10923..8992f4d8df6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -317,6 +317,19 @@ model LiteLLM_VerificationToken { @@index([budget_reset_at, expires]) } +// Deprecated keys during grace period - allows old key to work until revoke_at +model LiteLLM_DeprecatedVerificationToken { + id String @id @default(uuid()) + token String // Hashed old key + active_token_id String // Current token hash in LiteLLM_VerificationToken + revoke_at DateTime // When the old key stops working + created_at DateTime @default(now()) @map("created_at") + + @@unique([token]) + @@index([token, revoke_at]) + @@index([revoke_at]) +} + // Audit table for deleted keys - preserves spend and key information for historical tracking model LiteLLM_DeletedVerificationToken { id String @id @default(uuid()) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6bbf0df74de..c9b07ce447d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2644,6 +2644,71 @@ class PrismaClient: sql_query ) + # If not found in main table, check deprecated keys (grace period) + if response is None: + try: + deprecated_sql = f""" + SELECT active_token_id FROM "LiteLLM_DeprecatedVerificationToken" + WHERE token = '{hashed_token}' AND revoke_at > NOW() + LIMIT 1 + """ + deprecated_row = ( + await self._query_first_with_cached_plan_fallback( + deprecated_sql + ) + ) + if deprecated_row and deprecated_row.get("active_token_id"): + active_token_id = deprecated_row["active_token_id"] + sql_query_active = f""" + SELECT + v.*, + t.spend AS team_spend, + t.max_budget AS team_max_budget, + t.tpm_limit AS team_tpm_limit, + t.rpm_limit AS team_rpm_limit, + t.models AS team_models, + t.metadata AS team_metadata, + t.blocked AS team_blocked, + t.team_alias AS team_alias, + t.metadata AS team_metadata, + t.members_with_roles AS team_members_with_roles, + t.object_permission_id AS team_object_permission_id, + t.organization_id as org_id, + tm.spend AS team_member_spend, + m.aliases AS team_model_aliases, + b.max_budget AS litellm_budget_table_max_budget, + b.tpm_limit AS litellm_budget_table_tpm_limit, + b.rpm_limit AS litellm_budget_table_rpm_limit, + b.model_max_budget as litellm_budget_table_model_max_budget, + b.soft_budget as litellm_budget_table_soft_budget, + o.metadata as organization_metadata, + b2.max_budget as organization_max_budget, + b2.tpm_limit as organization_tpm_limit, + b2.rpm_limit as organization_rpm_limit + FROM "LiteLLM_VerificationToken" AS v + LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id + LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id + LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id + LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id + LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id + LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id + WHERE v.token = '{active_token_id}' + """ + response = ( + await self._query_first_with_cached_plan_fallback( + sql_query_active + ) + ) + if response is not None: + verbose_proxy_logger.debug( + "Deprecated key used during grace period" + ) + except Exception as deprecated_lookup_error: + verbose_proxy_logger.debug( + "Deprecated key lookup skipped: %s", + deprecated_lookup_error, + ) + if response is not None: if response["team_models"] is None: response["team_models"] = [] diff --git a/schema.prisma b/schema.prisma index 3b81da10923..8992f4d8df6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -317,6 +317,19 @@ model LiteLLM_VerificationToken { @@index([budget_reset_at, expires]) } +// Deprecated keys during grace period - allows old key to work until revoke_at +model LiteLLM_DeprecatedVerificationToken { + id String @id @default(uuid()) + token String // Hashed old key + active_token_id String // Current token hash in LiteLLM_VerificationToken + revoke_at DateTime // When the old key stops working + created_at DateTime @default(now()) @map("created_at") + + @@unique([token]) + @@index([token, revoke_at]) + @@index([revoke_at]) +} + // Audit table for deleted keys - preserves spend and key information for historical tracking model LiteLLM_DeletedVerificationToken { id String @id @default(uuid()) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py index 6b3b4c92416..1f5fab71ef3 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py @@ -4,7 +4,7 @@ Test key rotation manager functionality import os import sys from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest @@ -24,7 +24,7 @@ class TestKeyRotationManager: async def test_should_rotate_key_logic(self): """ Test the core logic for determining when a key should be rotated. - + This tests: - Keys with null key_rotation_at should rotate immediately - Keys with future key_rotation_at should not rotate @@ -33,69 +33,69 @@ class TestKeyRotationManager: # Setup mock_prisma_client = AsyncMock() manager = KeyRotationManager(mock_prisma_client) - + now = datetime.now(timezone.utc) - + # Test Case 1: No rotation time set (key_rotation_at = None) - should rotate key_no_rotation_time = LiteLLM_VerificationToken( token="test-token-1", auto_rotate=True, rotation_interval="30s", key_rotation_at=None, - rotation_count=0 + rotation_count=0, ) - - assert manager._should_rotate_key(key_no_rotation_time, now) == True - + + assert manager._should_rotate_key(key_no_rotation_time, now) is True + # Test Case 2: Future rotation time - should NOT rotate key_future_rotation = LiteLLM_VerificationToken( token="test-token-2", auto_rotate=True, rotation_interval="30s", key_rotation_at=now + timedelta(seconds=10), - rotation_count=1 + rotation_count=1, ) - - assert manager._should_rotate_key(key_future_rotation, now) == False - + + assert manager._should_rotate_key(key_future_rotation, now) is False + # Test Case 3: Past rotation time - should rotate key_past_rotation = LiteLLM_VerificationToken( token="test-token-3", auto_rotate=True, rotation_interval="30s", key_rotation_at=now - timedelta(seconds=10), - rotation_count=2 + rotation_count=2, ) - - assert manager._should_rotate_key(key_past_rotation, now) == True - + + assert manager._should_rotate_key(key_past_rotation, now) is True + # Test Case 4: Exact rotation time - should rotate key_exact_rotation = LiteLLM_VerificationToken( token="test-token-4", auto_rotate=True, rotation_interval="30s", key_rotation_at=now, - rotation_count=1 + rotation_count=1, ) - - assert manager._should_rotate_key(key_exact_rotation, now) == True - + + assert manager._should_rotate_key(key_exact_rotation, now) is True + # Test Case 5: No rotation interval - should NOT rotate key_no_interval = LiteLLM_VerificationToken( token="test-token-5", auto_rotate=True, rotation_interval=None, key_rotation_at=None, - rotation_count=0 + rotation_count=0, ) - - assert manager._should_rotate_key(key_no_interval, now) == False + + assert manager._should_rotate_key(key_no_interval, now) is False @pytest.mark.asyncio async def test_find_keys_needing_rotation(self): """ Test finding keys that need rotation from database. - + This tests: - Only keys with auto_rotate=True are considered - Database query filters by key_rotation_at properly @@ -104,10 +104,10 @@ class TestKeyRotationManager: # Setup mock_prisma_client = AsyncMock() manager = KeyRotationManager(mock_prisma_client) - + # Use a fixed timestamp to avoid timing issues in tests now = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) - + # Mock database response - these are the keys the database query would return mock_keys = [ LiteLLM_VerificationToken( @@ -115,42 +115,47 @@ class TestKeyRotationManager: auto_rotate=True, rotation_interval="30s", key_rotation_at=None, # Should rotate (null key_rotation_at) - rotation_count=0 + rotation_count=0, ), LiteLLM_VerificationToken( token="token-2", auto_rotate=True, rotation_interval="60s", - key_rotation_at=now - timedelta(seconds=10), # Should rotate (past time) - rotation_count=1 - ) + key_rotation_at=now + - timedelta(seconds=10), # Should rotate (past time) + rotation_count=1, + ), ] - - mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = mock_keys - + + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = ( + mock_keys + ) + # Mock datetime.now to return our fixed timestamp from unittest.mock import patch - with patch('litellm.proxy.common_utils.key_rotation_manager.datetime') as mock_datetime: + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.datetime" + ) as mock_datetime: mock_datetime.now.return_value = now - mock_datetime.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - + mock_datetime.side_effect = lambda *args, **kwargs: datetime( + *args, **kwargs + ) + # Execute keys_needing_rotation = await manager._find_keys_needing_rotation() - + # Verify database query - should use OR condition for key_rotation_at mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( where={ "auto_rotate": True, - "OR": [ - {"key_rotation_at": None}, - {"key_rotation_at": {"lte": now}} - ] + "OR": [{"key_rotation_at": None}, {"key_rotation_at": {"lte": now}}], } ) - + # Verify all keys returned by database query are included (no additional filtering) assert len(keys_needing_rotation) == 2 - + tokens_needing_rotation = [key.token for key in keys_needing_rotation] assert "token-1" in tokens_needing_rotation # Null key_rotation_at assert "token-2" in tokens_needing_rotation # Past key_rotation_at @@ -159,7 +164,7 @@ class TestKeyRotationManager: async def test_rotate_key_updates_database(self): """ Test that key rotation properly updates the database with new rotation info. - + This tests: - Rotation count is incremented - last_rotation_at is set to current time @@ -169,7 +174,7 @@ class TestKeyRotationManager: # Setup mock_prisma_client = AsyncMock() manager = KeyRotationManager(mock_prisma_client) - + # Mock key to rotate key_to_rotate = LiteLLM_VerificationToken( token="old-token", @@ -177,31 +182,35 @@ class TestKeyRotationManager: rotation_interval="30s", last_rotation_at=None, key_rotation_at=None, - rotation_count=0 + rotation_count=0, ) - + # Mock regenerate_key_fn response mock_response = GenerateKeyResponse( - key="new-api-key", - token_id="new-token-id", - user_id="test-user" + key="new-api-key", token_id="new-token-id", user_id="test-user" ) - + # Mock the regenerate function from unittest.mock import patch - with patch('litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn', return_value=mock_response): - with patch('litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook'): + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook" + ): # Execute await manager._rotate_key(key_to_rotate) - + # Verify database update was called with correct data mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once() - + call_args = mock_prisma_client.db.litellm_verificationtoken.update.call_args - + # Check the WHERE clause targets the new token assert call_args[1]["where"]["token"] == "new-token-id" - + # Check the data being updated update_data = call_args[1]["data"] assert update_data["rotation_count"] == 1 # Incremented from 0 @@ -209,9 +218,75 @@ class TestKeyRotationManager: assert isinstance(update_data["last_rotation_at"], datetime) assert "key_rotation_at" in update_data assert isinstance(update_data["key_rotation_at"], datetime) - + # Verify key_rotation_at is set to future time (30s from now) now = datetime.now(timezone.utc) next_rotation = update_data["key_rotation_at"] time_diff = (next_rotation - now).total_seconds() - assert 25 <= time_diff <= 35 # Should be around 30 seconds, allow some tolerance + assert ( + 25 <= time_diff <= 35 + ) # Should be around 30 seconds, allow some tolerance + + @pytest.mark.asyncio + async def test_cleanup_expired_deprecated_keys(self): + """ + Test that _cleanup_expired_deprecated_keys deletes expired deprecated keys. + """ + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.return_value = { + "count": 3 + } + manager = KeyRotationManager(mock_prisma_client) + + await manager._cleanup_expired_deprecated_keys() + + mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.assert_called_once() + call_args = ( + mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.call_args + ) + assert "revoke_at" in call_args[1]["where"] + assert call_args[1]["where"]["revoke_at"]["lt"] is not None + + @pytest.mark.asyncio + async def test_rotate_key_passes_grace_period_hours(self): + """ + Test that _rotate_key passes grace_period_hours in RegenerateKeyRequest. + """ + mock_prisma_client = AsyncMock() + manager = KeyRotationManager(mock_prisma_client) + + key_to_rotate = LiteLLM_VerificationToken( + token="old-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + mock_response = GenerateKeyResponse( + key="new-api-key", + token_id="new-token-id", + user_id="test-user", + ) + + from unittest.mock import patch + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + ) as mock_regenerate: + mock_regenerate.return_value = mock_response + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS", + 48, + ): + await manager._rotate_key(key_to_rotate) + + mock_regenerate.assert_called_once() + call_args = mock_regenerate.call_args + regenerate_request = call_args[1]["data"] + assert regenerate_request.grace_period_hours == 48 diff --git a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx index a4339e11920..3cea198c62a 100644 --- a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx @@ -37,6 +37,7 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat tpm_limit: selectedToken.tpm_limit, rpm_limit: selectedToken.rpm_limit, duration: selectedToken.duration || "", + grace_period_hours: 0, }); // Initialize the current access token @@ -223,6 +224,17 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat Current expiry: {selectedToken?.expires ? new Date(selectedToken.expires).toLocaleString() : "Never"} {newExpiryTime &&
New expiry: {newExpiryTime}
} + + + +
+ Recommended: 24-72 hours for production keys to allow seamless client migration. +
)} From 65401138b5856c095b83cb7f0011cdfc352d30a4 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Tue, 3 Feb 2026 10:49:49 +0530 Subject: [PATCH 03/50] fix: refactor, race condition handle, fstring sql injection --- .../key_management_endpoints.py | 19 ++++-- litellm/proxy/utils.py | 65 ++++--------------- 2 files changed, 28 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index dea1e55f7f9..55f96e4189c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3365,12 +3365,21 @@ async def regenerate_key_fn( # noqa: PLR0915 revoke_at = datetime.now(timezone.utc) + timedelta( hours=grace_period_hours ) - await prisma_client.db.litellm_deprecatedverificationtoken.create( + # Use upsert to handle concurrent rotations gracefully; avoids + # unique constraint violation if same key is rotated simultaneously + await prisma_client.db.litellm_deprecatedverificationtoken.upsert( + where={"token": hashed_api_key}, data={ - "token": hashed_api_key, - "active_token_id": new_token_hash, - "revoke_at": revoke_at, - } + "create": { + "token": hashed_api_key, + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + }, + "update": { + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + }, + }, ) verbose_proxy_logger.debug( "Deprecated key retained for %s hours (revoke_at: %s)", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c9b07ce447d..5cc5507dd09 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7,7 +7,7 @@ import smtplib import threading import time import traceback -from datetime import date, datetime, timedelta +from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from typing import ( @@ -2647,57 +2647,20 @@ class PrismaClient: # If not found in main table, check deprecated keys (grace period) if response is None: try: - deprecated_sql = f""" - SELECT active_token_id FROM "LiteLLM_DeprecatedVerificationToken" - WHERE token = '{hashed_token}' AND revoke_at > NOW() - LIMIT 1 - """ - deprecated_row = ( - await self._query_first_with_cached_plan_fallback( - deprecated_sql - ) + deprecated_row = await self.db.litellm_deprecatedverificationtoken.find_first( + where={ + "token": hashed_token, + "revoke_at": {"gt": datetime.now(timezone.utc)}, + }, + select={"active_token_id": True}, ) - if deprecated_row and deprecated_row.get("active_token_id"): - active_token_id = deprecated_row["active_token_id"] - sql_query_active = f""" - SELECT - v.*, - t.spend AS team_spend, - t.max_budget AS team_max_budget, - t.tpm_limit AS team_tpm_limit, - t.rpm_limit AS team_rpm_limit, - t.models AS team_models, - t.metadata AS team_metadata, - t.blocked AS team_blocked, - t.team_alias AS team_alias, - t.metadata AS team_metadata, - t.members_with_roles AS team_members_with_roles, - t.object_permission_id AS team_object_permission_id, - t.organization_id as org_id, - tm.spend AS team_member_spend, - m.aliases AS team_model_aliases, - b.max_budget AS litellm_budget_table_max_budget, - b.tpm_limit AS litellm_budget_table_tpm_limit, - b.rpm_limit AS litellm_budget_table_rpm_limit, - b.model_max_budget as litellm_budget_table_model_max_budget, - b.soft_budget as litellm_budget_table_soft_budget, - o.metadata as organization_metadata, - b2.max_budget as organization_max_budget, - b2.tpm_limit as organization_tpm_limit, - b2.rpm_limit as organization_rpm_limit - FROM "LiteLLM_VerificationToken" AS v - LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id - LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id - LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id - LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id - LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id - LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id - WHERE v.token = '{active_token_id}' - """ - response = ( - await self._query_first_with_cached_plan_fallback( - sql_query_active - ) + if deprecated_row and deprecated_row.active_token_id: + response = await self.get_data( + token=deprecated_row.active_token_id, + table_name="combined_view", + query_type="find_unique", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) if response is not None: verbose_proxy_logger.debug( From 1792b3c8e5b61f4e5a9951fb53e31542d40c5490 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Fri, 13 Feb 2026 03:28:51 +0530 Subject: [PATCH 04/50] fix: add async call to avoid server pauses --- litellm/proxy/management_endpoints/ui_sso.py | 16 ++++++++-------- .../proxy/management_endpoints/test_ui_sso.py | 19 +++++++++++-------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a2d7510c344..2760197781b 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -695,7 +695,7 @@ async def get_generic_sso_response( try: result = await generic_sso.verify_and_process( request, - params=SSOAuthenticationHandler.prepare_token_exchange_parameters( + params=await SSOAuthenticationHandler.prepare_token_exchange_parameters( request=request, generic_include_client_id=generic_include_client_id, ), @@ -1721,13 +1721,13 @@ class SSOAuthenticationHandler: # so callbacks landing on another pod can retrieve it (multi-pod SSO). cache_key = f"pkce_verifier:{redirect_params['state']}" if redis_usage_cache is not None: - redis_usage_cache.set_cache( + await redis_usage_cache.async_set_cache( key=cache_key, value=json.dumps(code_verifier), ttl=600, ) else: - user_api_key_cache.set_cache( + await user_api_key_cache.async_set_cache( key=cache_key, value=code_verifier, ttl=600, @@ -2328,7 +2328,7 @@ class SSOAuthenticationHandler: return redirect_response @staticmethod - def prepare_token_exchange_parameters( + async def prepare_token_exchange_parameters( request: Request, generic_include_client_id: bool, ) -> dict: @@ -2354,9 +2354,9 @@ class SSOAuthenticationHandler: cache_key = f"pkce_verifier:{state}" if redis_usage_cache is not None: - code_verifier = redis_usage_cache.get_cache(key=cache_key) + code_verifier = await redis_usage_cache.async_get_cache(key=cache_key) else: - code_verifier = user_api_key_cache.get_cache(key=cache_key) + code_verifier = await user_api_key_cache.async_get_cache(key=cache_key) if code_verifier: # Add code_verifier to token exchange parameters (Redis returns decoded string) @@ -2371,9 +2371,9 @@ class SSOAuthenticationHandler: # Clean up the cache entry (single-use verifier) if redis_usage_cache is not None: - redis_usage_cache.delete_cache(key=cache_key) + await redis_usage_cache.async_delete_cache(key=cache_key) else: - user_api_key_cache.delete_cache(key=cache_key) + await user_api_key_cache.async_delete_cache(key=cache_key) return token_params @staticmethod 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 3a5505cdbc5..1a67f7f6b6f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3073,14 +3073,15 @@ class TestPKCEFunctionality: test_state = "test_oauth_state_123" mock_request.query_params = {"state": test_state} - # Mock cache + # Mock cache with async methods mock_cache = MagicMock() test_code_verifier = "test_code_verifier_abc123xyz" - mock_cache.get_cache.return_value = test_code_verifier + mock_cache.async_get_cache = AsyncMock(return_value=test_code_verifier) + mock_cache.async_delete_cache = AsyncMock() with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): # Act - token_params = SSOAuthenticationHandler.prepare_token_exchange_parameters( + token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( request=mock_request, generic_include_client_id=False ) @@ -3089,10 +3090,10 @@ class TestPKCEFunctionality: assert token_params["code_verifier"] == test_code_verifier # Verify cache was accessed and deleted - mock_cache.get_cache.assert_called_once_with( + mock_cache.async_get_cache.assert_called_once_with( key=f"pkce_verifier:{test_state}" ) - mock_cache.delete_cache.assert_called_once_with( + mock_cache.async_delete_cache.assert_called_once_with( key=f"pkce_verifier:{test_state}" ) @@ -3117,6 +3118,8 @@ class TestPKCEFunctionality: test_state = "test456" mock_cache = MagicMock() + mock_cache.async_set_cache = AsyncMock() + with patch.dict(os.environ, {"GENERIC_CLIENT_USE_PKCE": "true"}): with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): # Act @@ -3127,9 +3130,9 @@ class TestPKCEFunctionality: ) # Assert - # Verify cache was called to store code_verifier - mock_cache.set_cache.assert_called_once() - cache_call = mock_cache.set_cache.call_args + # Verify async cache was called to store code_verifier + mock_cache.async_set_cache.assert_called_once() + cache_call = mock_cache.async_set_cache.call_args assert cache_call.kwargs["key"] == f"pkce_verifier:{test_state}" assert cache_call.kwargs["ttl"] == 600 assert len(cache_call.kwargs["value"]) == 43 From bc5543cfdcbe2dd3318dd3e126e478efd3e833ec Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Fri, 13 Feb 2026 03:39:57 +0530 Subject: [PATCH 05/50] Update tests/test_litellm/proxy/management_endpoints/test_ui_sso.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 1a67f7f6b6f..d3cbe6e1718 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3209,7 +3209,7 @@ class TestPKCEFunctionality: mock_request = MagicMock(spec=Request) mock_request.query_params = {"state": "multi_pod_state_xyz"} token_params = ( - SSOAuthenticationHandler.prepare_token_exchange_parameters( + await SSOAuthenticationHandler.prepare_token_exchange_parameters( request=mock_request, generic_include_client_id=False ) ) From eb249b2f06a9c74a06874874838e3a0e97b5a5ea Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Fri, 13 Feb 2026 03:46:31 +0530 Subject: [PATCH 06/50] fix: add await in tests --- .../proxy/management_endpoints/test_ui_sso.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) 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 1a67f7f6b6f..bd1470a6d65 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3081,8 +3081,10 @@ class TestPKCEFunctionality: with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): # Act - token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False + token_params = ( + await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False + ) ) # Assert @@ -3208,10 +3210,8 @@ class TestPKCEFunctionality: # Pod B: callback with same state, retrieve from "Redis" mock_request = MagicMock(spec=Request) mock_request.query_params = {"state": "multi_pod_state_xyz"} - token_params = ( - SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False - ) + token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False ) assert "code_verifier" in token_params assert token_params["code_verifier"] == json.loads(stored_value) @@ -3277,10 +3277,8 @@ class TestPKCEFunctionality: # Same pod: callback retrieves from in-memory cache mock_request = MagicMock(spec=Request) mock_request.query_params = {"state": "fallback_state_xyz"} - token_params = ( - SSOAuthenticationHandler.prepare_token_exchange_parameters( - request=mock_request, generic_include_client_id=False - ) + token_params = await SSOAuthenticationHandler.prepare_token_exchange_parameters( + request=mock_request, generic_include_client_id=False ) assert "code_verifier" in token_params assert token_params["code_verifier"] == stored_value @@ -3306,7 +3304,7 @@ class TestPKCEFunctionality: mock_request = MagicMock(spec=Request) mock_request.query_params = {} token_params = ( - SSOAuthenticationHandler.prepare_token_exchange_parameters( + await SSOAuthenticationHandler.prepare_token_exchange_parameters( request=mock_request, generic_include_client_id=False ) ) From e0846389e9ec249ffb257f6240370ad984356204 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Fri, 13 Feb 2026 04:04:03 +0530 Subject: [PATCH 07/50] add modify test to perform async run --- litellm/proxy/management_endpoints/ui_sso.py | 3 +- .../proxy/management_endpoints/test_ui_sso.py | 42 +++++++++++-------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 2760197781b..1dc239e50a4 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -11,7 +11,6 @@ Has all /sso/* routes import asyncio import base64 import hashlib -import json import os import secrets from copy import deepcopy @@ -1723,7 +1722,7 @@ class SSOAuthenticationHandler: if redis_usage_cache is not None: await redis_usage_cache.async_set_cache( key=cache_key, - value=json.dumps(code_verifier), + value=code_verifier, ttl=600, ) else: 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 bd1470a6d65..44aa04c0ad6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3159,10 +3159,10 @@ class TestPKCEFunctionality: def __init__(self): self._store = {} - def set_cache(self, key, value, **kwargs): + async def async_set_cache(self, key, value, **kwargs): self._store[key] = value - def get_cache(self, key, **kwargs): + async def async_get_cache(self, key, **kwargs): val = self._store.get(key) if val is None: return None @@ -3174,7 +3174,7 @@ class TestPKCEFunctionality: return val return val - def delete_cache(self, key): + async def async_delete_cache(self, key): self._store.pop(key, None) mock_redis = MockRedisCache() @@ -3200,7 +3200,7 @@ class TestPKCEFunctionality: state="multi_pod_state_xyz", generic_authorization_endpoint="https://auth.example.com/authorize", ) - mock_in_memory.set_cache.assert_not_called() + mock_in_memory.async_set_cache.assert_not_called() # MockRedisCache is a real class; assert on state, not .assert_called_* stored_key = "pkce_verifier:multi_pod_state_xyz" assert stored_key in mock_redis._store @@ -3215,7 +3215,7 @@ class TestPKCEFunctionality: ) assert "code_verifier" in token_params assert token_params["code_verifier"] == json.loads(stored_value) - mock_in_memory.get_cache.assert_not_called() + mock_in_memory.async_get_cache.assert_not_called() # delete_cache called; key removed (asserted below) # Verifier consumed (single-use); key removed from "Redis" @@ -3234,19 +3234,19 @@ class TestPKCEFunctionality: # In-memory store (simulates user_api_key_cache on one pod) in_memory_store = {} - def set_cache(key, value, **kwargs): + async def async_set_cache(key, value, **kwargs): in_memory_store[key] = value - def get_cache(key, **kwargs): + async def async_get_cache(key, **kwargs): return in_memory_store.get(key) - def delete_cache(key): + async def async_delete_cache(key): in_memory_store.pop(key, None) mock_in_memory = MagicMock() - mock_in_memory.set_cache.side_effect = set_cache - mock_in_memory.get_cache.side_effect = get_cache - mock_in_memory.delete_cache.side_effect = delete_cache + mock_in_memory.async_set_cache = AsyncMock(side_effect=async_set_cache) + mock_in_memory.async_get_cache = AsyncMock(side_effect=async_get_cache) + mock_in_memory.async_delete_cache = AsyncMock(side_effect=async_delete_cache) mock_sso = MagicMock() mock_redirect_response = MagicMock() @@ -3268,9 +3268,11 @@ class TestPKCEFunctionality: state="fallback_state_xyz", generic_authorization_endpoint="https://auth.example.com/authorize", ) - mock_in_memory.set_cache.assert_called_once() - stored_key = mock_in_memory.set_cache.call_args.kwargs["key"] - stored_value = mock_in_memory.set_cache.call_args.kwargs["value"] + mock_in_memory.async_set_cache.assert_called_once() + stored_key = mock_in_memory.async_set_cache.call_args.kwargs["key"] + stored_value = mock_in_memory.async_set_cache.call_args.kwargs[ + "value" + ] assert stored_key == "pkce_verifier:fallback_state_xyz" assert isinstance(stored_value, str) and len(stored_value) == 43 @@ -3282,8 +3284,12 @@ class TestPKCEFunctionality: ) assert "code_verifier" in token_params assert token_params["code_verifier"] == stored_value - mock_in_memory.get_cache.assert_called_once_with(key=stored_key) - mock_in_memory.delete_cache.assert_called_once_with(key=stored_key) + mock_in_memory.async_get_cache.assert_called_once_with( + key=stored_key + ) + mock_in_memory.async_delete_cache.assert_called_once_with( + key=stored_key + ) # Verifier consumed; key removed from in-memory assert "pkce_verifier:fallback_state_xyz" not in in_memory_store @@ -3309,8 +3315,8 @@ class TestPKCEFunctionality: ) ) assert "code_verifier" not in token_params - mock_redis.get_cache.assert_not_called() - mock_in_memory.get_cache.assert_not_called() + mock_redis.async_get_cache.assert_not_called() + mock_in_memory.async_get_cache.assert_not_called() # Tests for SSO user team assignment bug (Issue: SSO Users Not Added to Entra-Synced Teams on First Login) From a2b4728e749ada2c0d315fff03602ae69525cfd0 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Fri, 13 Feb 2026 04:40:29 +0530 Subject: [PATCH 08/50] Update tests/test_litellm/proxy/management_endpoints/test_ui_sso.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 44aa04c0ad6..d519ce43eee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3205,7 +3205,7 @@ class TestPKCEFunctionality: stored_key = "pkce_verifier:multi_pod_state_xyz" assert stored_key in mock_redis._store stored_value = mock_redis._store[stored_key] - assert json.loads(stored_value) # valid verifier string + assert isinstance(stored_value, str) and len(stored_value) == 43 # Pod B: callback with same state, retrieve from "Redis" mock_request = MagicMock(spec=Request) From c56bbb906721b0a348e2cff295ec228d3462e5e8 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Fri, 13 Feb 2026 04:48:36 +0530 Subject: [PATCH 09/50] Update tests/test_litellm/proxy/management_endpoints/test_ui_sso.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d519ce43eee..016131f24dc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3160,7 +3160,7 @@ class TestPKCEFunctionality: self._store = {} async def async_set_cache(self, key, value, **kwargs): - self._store[key] = value + self._store[key] = json.dumps(value) async def async_get_cache(self, key, **kwargs): val = self._store.get(key) From f87b12a2f7d2596d8b85161b7157b55e0981513b Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Sat, 14 Feb 2026 04:18:08 +0530 Subject: [PATCH 10/50] fix grace period with better error handling on frontend and as per best practices --- docs/my-website/docs/proxy/config_settings.md | 2 +- docs/my-website/docs/proxy/virtual_keys.md | 8 +- litellm/constants.py | 18 ++- litellm/proxy/_types.py | 90 ++++++------- .../common_utils/key_rotation_manager.py | 4 +- .../key_management_endpoints.py | 118 ++++++++++++------ litellm/proxy/utils.py | 99 ++++++++++----- .../common_utils/test_key_rotation_manager.py | 10 +- .../organisms/regenerate_key_modal.tsx | 18 ++- 9 files changed, 222 insertions(+), 145 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 5499e541335..03cf665b3c8 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -765,7 +765,7 @@ router_settings: | LITELM_ENVIRONMENT | Environment of LiteLLM Instance, used by logging services. Currently only used by DeepEval. | LITELLM_KEY_ROTATION_ENABLED | Enable auto-key rotation for LiteLLM (boolean). Default is false. | LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS | Interval in seconds for how often to run job that auto-rotates keys. Default is 86400 (24 hours). -| LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS | Hours to keep old key valid after rotation (e.g. 24, 48, 72). Default is 0 (immediate revoke). Used for scheduled rotations and as fallback when not specified in regenerate request. +| LITELLM_KEY_ROTATION_GRACE_PERIOD | Duration to keep old key valid after rotation (e.g. "24h", "2d"). Default is empty (immediate revoke). Used for scheduled rotations and as fallback when not specified in regenerate request. | LITELLM_LICENSE | License key for LiteLLM usage | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM | LITELLM_LOG | Enable detailed logging for LiteLLM diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index 31506ea2611..c74aa75ff4a 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -550,12 +550,12 @@ curl 'http://localhost:4000/key/sk-1234/regenerate' \ "gpt-4", "gpt-3.5-turbo" ], - "grace_period_hours": 48 + "grace_period": "48h" }' ``` -**Grace period (optional)**: Set `grace_period_hours` (e.g. 24, 48, 72) to keep the old key valid for a transitional period. Both old and new keys work until the grace period elapses, enabling seamless cutover without production downtime. Default is 0 (immediate revoke). Can also be set via `LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS` for scheduled rotations. +**Grace period (optional)**: Set `grace_period` (e.g. `"24h"`, `"2d"`, `"1w"`) to keep the old key valid for a transitional period. Both old and new keys work until the grace period elapses, enabling seamless cutover without production downtime. Omitted or empty = immediate revoke. Can also be set via `LITELLM_KEY_ROTATION_GRACE_PERIOD` env var for scheduled rotations. **Read More** @@ -643,13 +643,13 @@ Set these environment variables when starting the proxy: |----------|-------------|---------| | `LITELLM_KEY_ROTATION_ENABLED` | Enable the rotation worker | `false` | | `LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS` | How often to scan for keys to rotate (in seconds) | `86400` (24 hours) | -| `LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS` | Hours to keep old key valid after rotation (24–72 recommended) | `0` (immediate revoke) | +| `LITELLM_KEY_ROTATION_GRACE_PERIOD` | Duration to keep old key valid after rotation (e.g. `24h`, `2d`) | `""` (immediate revoke) | **Example:** ```bash export LITELLM_KEY_ROTATION_ENABLED=true export LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS=3600 # Check every hour -export LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS=48 # Keep old key valid for 48h during cutover +export LITELLM_KEY_ROTATION_GRACE_PERIOD=48h # Keep old key valid for 48h during cutover litellm --config config.yaml ``` diff --git a/litellm/constants.py b/litellm/constants.py index 73713e6a05b..56349505299 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -106,9 +106,7 @@ MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int( # npm/npx needs a writable cache dir; in containers the default (~/.npm) # may not exist or be read-only. /tmp is always writable. MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") -MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int( - os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10") -) +MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) LITELLM_UI_ALLOW_HEADERS = [ "x-litellm-semantic-filter", @@ -225,9 +223,7 @@ REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 2000)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth -LITELLM_ASYNCIO_QUEUE_MAXSIZE = int( - os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000) -) +LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int( os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000) ) @@ -349,7 +345,9 @@ MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int( DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000)) #### Networking settings #### request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", 6000)) # time in seconds -DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes +DEFAULT_A2A_AGENT_TIMEOUT: float = float( + os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000) +) # 10 minutes # Patterns that indicate a localhost/internal URL in A2A agent cards that should be # replaced with the original base_url. This is a common misconfiguration where # developers deploy agents with development URLs in their agent cards. @@ -1256,9 +1254,9 @@ LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false" LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int( os.getenv("LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400) ) # 24 hours default -LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS = int( - os.getenv("LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS", 0) -) # Hours to keep old key valid after rotation; 0 = immediate revoke (default) +LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv( + "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" +) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default) UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b8c84c7fa2d..fb632dd16d4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -854,9 +854,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -995,9 +995,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period_hours: Optional[ - int - ] = None # Hours to keep old key valid; 0/None = immediate revoke + grace_period: Optional[ + str + ] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1409,12 +1409,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model @model_validator(mode="before") @classmethod @@ -1436,12 +1436,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -1530,15 +1530,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1630,9 +1630,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1964,9 +1964,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2406,9 +2406,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None model_config = ConfigDict(protected_namespaces=()) @@ -3384,9 +3384,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -3604,9 +3604,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -3749,9 +3749,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 7ef26bf07b9..88fd08f6ab4 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -10,7 +10,7 @@ from typing import List from litellm._logging import verbose_proxy_logger from litellm.constants import ( LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, - LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS, + LITELLM_KEY_ROTATION_GRACE_PERIOD, ) from litellm.proxy._types import ( GenerateKeyResponse, @@ -143,7 +143,7 @@ class KeyRotationManager: regenerate_request = RegenerateKeyRequest( key=key.token or "", key_alias=key.key_alias, # Pass key alias to ensure correct secret is updated in AWS Secrets Manager - grace_period_hours=LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS, + grace_period=LITELLM_KEY_ROTATION_GRACE_PERIOD or None, ) # Create a system user for key rotation diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f7ad44879bd..e8a1af27978 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -629,7 +629,11 @@ async def _common_key_generation_helper( # noqa: PLR0915 # Validate user-provided key format if data.key is not None and not data.key.startswith("sk-"): - _masked = "{}****{}".format(data.key[:4], data.key[-4:]) if len(data.key) > 8 else "****" + _masked = ( + "{}****{}".format(data.key[:4], data.key[-4:]) + if len(data.key) > 8 + else "****" + ) raise HTTPException( status_code=400, detail={ @@ -1342,7 +1346,7 @@ async def prepare_key_update_data( data_json: dict = data.model_dump(exclude_unset=True) data_json.pop("key", None) data_json.pop("new_key", None) - data_json.pop("grace_period_hours", None) # Request-only param, not a DB column + data_json.pop("grace_period", None) # Request-only param, not a DB column if ( data.metadata is not None and data.metadata.get("service_account_id") is not None @@ -3178,6 +3182,69 @@ def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: return new_token +async def _insert_deprecated_key( + prisma_client: "PrismaClient", + old_token_hash: str, + new_token_hash: str, + grace_period: Optional[str], +) -> None: + """ + Insert old key into deprecated table so it remains valid during grace period. + + Uses upsert to handle concurrent rotations gracefully. + + Parameters: + prisma_client: DB client + old_token_hash: Hash of the old key being rotated out + new_token_hash: Hash of the new replacement key + grace_period: Duration string (e.g. "24h", "2d") or None/empty for immediate revoke + """ + grace_period_value = grace_period or os.getenv( + "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" + ) + if not grace_period_value: + return + + try: + grace_seconds = duration_in_seconds(grace_period_value) + except ValueError: + verbose_proxy_logger.warning( + "Invalid grace_period format: %s. Expected format like '24h', '2d'.", + grace_period_value, + ) + return + + if grace_seconds <= 0: + return + + try: + revoke_at = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) + await prisma_client.db.litellm_deprecatedverificationtoken.upsert( + where={"token": old_token_hash}, + data={ + "create": { + "token": old_token_hash, + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + }, + "update": { + "active_token_id": new_token_hash, + "revoke_at": revoke_at, + }, + }, + ) + verbose_proxy_logger.debug( + "Deprecated key retained for %s (revoke_at: %s)", + grace_period_value, + revoke_at, + ) + except Exception as deprecated_err: + verbose_proxy_logger.warning( + "Failed to insert deprecated key for grace period: %s", + deprecated_err, + ) + + @router.post( "/key/{key:path}/regenerate", tags=["key management"], @@ -3228,7 +3295,7 @@ async def regenerate_key_fn( # noqa: PLR0915 - permissions: Optional[dict] - Key-specific permissions - guardrails: Optional[List[str]] - List of active guardrails for the key - blocked: Optional[bool] - Whether the key is blocked - - grace_period_hours: Optional[int] - Hours to keep old key valid after rotation (e.g. 24, 48, 72). 0 or omitted = immediate revoke. Env: LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS + - grace_period: Optional[str] - Duration to keep old key valid after rotation (e.g. "24h", "2d"). Omitted = immediate revoke. Env: LITELLM_KEY_ROTATION_GRACE_PERIOD Returns: @@ -3369,44 +3436,13 @@ async def regenerate_key_fn( # noqa: PLR0915 update_data.update(non_default_values) update_data = prisma_client.jsonify_object(data=update_data) - # If grace period > 0, insert deprecated key so old key remains valid - if data is not None and data.grace_period_hours is not None: - grace_period_hours = data.grace_period_hours - else: - grace_period_hours = int( - os.getenv("LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS", "0") - ) - if grace_period_hours > 0: - try: - revoke_at = datetime.now(timezone.utc) + timedelta( - hours=grace_period_hours - ) - # Use upsert to handle concurrent rotations gracefully; avoids - # unique constraint violation if same key is rotated simultaneously - await prisma_client.db.litellm_deprecatedverificationtoken.upsert( - where={"token": hashed_api_key}, - data={ - "create": { - "token": hashed_api_key, - "active_token_id": new_token_hash, - "revoke_at": revoke_at, - }, - "update": { - "active_token_id": new_token_hash, - "revoke_at": revoke_at, - }, - }, - ) - verbose_proxy_logger.debug( - "Deprecated key retained for %s hours (revoke_at: %s)", - grace_period_hours, - revoke_at, - ) - except Exception as deprecated_err: - verbose_proxy_logger.warning( - "Failed to insert deprecated key for grace period: %s", - deprecated_err, - ) + # If grace period set, insert deprecated key so old key remains valid + await _insert_deprecated_key( + prisma_client=prisma_client, + old_token_hash=hashed_api_key, + new_token_hash=new_token_hash, + grace_period=data.grace_period if data else None, + ) # Update the token in the database updated_token = await prisma_client.db.litellm_verificationtoken.update( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5fbc003ceb1..528415c891d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1375,11 +1375,11 @@ class ProxyLogging: # Note: user_info is a CallInfo that can represent user/team/org level info. For team budgets, # alert_emails is populated from team_object.metadata.soft_budget_alerting_emails (see auth_checks.py) is_soft_budget_with_alert_emails = ( - type == "soft_budget" - and user_info.alert_emails is not None + type == "soft_budget" + and user_info.alert_emails is not None and len(user_info.alert_emails) > 0 ) - + if self.alerting is None and not is_soft_budget_with_alert_emails: # do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails) return @@ -1395,10 +1395,9 @@ class ProxyLogging: # 1. "email" is in alerting config, OR # 2. It's a soft_budget alert with team-specific alert_emails (bypasses global alerting config) should_send_email = ( - (self.alerting is not None and "email" in self.alerting) - or is_soft_budget_with_alert_emails - ) - + self.alerting is not None and "email" in self.alerting + ) or is_soft_budget_with_alert_emails + if should_send_email and self.email_logging_instance is not None: await self.email_logging_instance.budget_alerts( type=type, @@ -2032,6 +2031,54 @@ def jsonify_object(data: dict) -> dict: return db_data +# In-memory cache for deprecated key lookups: maps old_token_hash -> (active_token_id, expires_at_ts) +# Avoids a DB query on every auth request for non-deprecated keys. +_deprecated_key_cache: Dict[str, tuple] = {} +_DEPRECATED_KEY_CACHE_TTL_SECONDS = 60 + + +async def _lookup_deprecated_key( + db: Any, + hashed_token: str, +) -> Optional[str]: + """ + Check if a token exists in the deprecated keys table and is still within its grace period. + + Returns the active_token_id if found and valid, otherwise None. + Uses an in-memory cache to avoid DB queries on every auth request. + """ + now = datetime.now(timezone.utc) + now_ts = now.timestamp() + + # Check cache first + cached = _deprecated_key_cache.get(hashed_token) + if cached is not None: + active_token_id, cache_expires_at_ts = cached + if now_ts < cache_expires_at_ts: + return active_token_id + else: + _deprecated_key_cache.pop(hashed_token, None) + + try: + deprecated_row = await db.litellm_deprecatedverificationtoken.find_first( + where={ + "token": hashed_token, + "revoke_at": {"gt": now}, + }, + select={"active_token_id": True}, + ) + if deprecated_row and deprecated_row.active_token_id: + _deprecated_key_cache[hashed_token] = ( + deprecated_row.active_token_id, + now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, + ) + return deprecated_row.active_token_id + except Exception as e: + verbose_proxy_logger.debug("Deprecated key lookup skipped: %s", e) + + return None + + class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() @@ -2665,32 +2712,22 @@ class PrismaClient: ) # If not found in main table, check deprecated keys (grace period) - if response is None: - try: - deprecated_row = await self.db.litellm_deprecatedverificationtoken.find_first( - where={ - "token": hashed_token, - "revoke_at": {"gt": datetime.now(timezone.utc)}, - }, - select={"active_token_id": True}, + if response is None and hashed_token is not None: + active_token_id = await _lookup_deprecated_key( + db=self.db, hashed_token=hashed_token + ) + if active_token_id: + response = await self.get_data( + token=active_token_id, + table_name="combined_view", + query_type="find_unique", + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - if deprecated_row and deprecated_row.active_token_id: - response = await self.get_data( - token=deprecated_row.active_token_id, - table_name="combined_view", - query_type="find_unique", - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + if response is not None: + verbose_proxy_logger.debug( + "Deprecated key used during grace period" ) - if response is not None: - verbose_proxy_logger.debug( - "Deprecated key used during grace period" - ) - except Exception as deprecated_lookup_error: - verbose_proxy_logger.debug( - "Deprecated key lookup skipped: %s", - deprecated_lookup_error, - ) if response is not None: if response["team_models"] is None: diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py index 1f5fab71ef3..002b015d960 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py @@ -248,9 +248,9 @@ class TestKeyRotationManager: assert call_args[1]["where"]["revoke_at"]["lt"] is not None @pytest.mark.asyncio - async def test_rotate_key_passes_grace_period_hours(self): + async def test_rotate_key_passes_grace_period(self): """ - Test that _rotate_key passes grace_period_hours in RegenerateKeyRequest. + Test that _rotate_key passes grace_period in RegenerateKeyRequest. """ mock_prisma_client = AsyncMock() manager = KeyRotationManager(mock_prisma_client) @@ -281,12 +281,12 @@ class TestKeyRotationManager: new_callable=AsyncMock, ): with patch( - "litellm.proxy.common_utils.key_rotation_manager.LITELLM_KEY_ROTATION_GRACE_PERIOD_HOURS", - 48, + "litellm.proxy.common_utils.key_rotation_manager.LITELLM_KEY_ROTATION_GRACE_PERIOD", + "48h", ): await manager._rotate_key(key_to_rotate) mock_regenerate.assert_called_once() call_args = mock_regenerate.call_args regenerate_request = call_args[1]["data"] - assert regenerate_request.grace_period_hours == 48 + assert regenerate_request.grace_period == "48h" diff --git a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx index 3cea198c62a..2fad101c20f 100644 --- a/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx +++ b/ui/litellm-dashboard/src/components/organisms/regenerate_key_modal.tsx @@ -37,7 +37,7 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat tpm_limit: selectedToken.tpm_limit, rpm_limit: selectedToken.rpm_limit, duration: selectedToken.duration || "", - grace_period_hours: 0, + grace_period: "", }); // Initialize the current access token @@ -225,15 +225,21 @@ export function RegenerateKeyModal({ selectedToken, visible, onClose, onKeyUpdat {newExpiryTime &&
New expiry: {newExpiryTime}
} - +
- Recommended: 24-72 hours for production keys to allow seamless client migration. + Recommended: 24h to 72h for production keys to allow seamless client migration.
)} From 37fc4a35ff9b33babfcdd17bc6abe06fa4f508ca Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sat, 14 Feb 2026 04:27:13 +0530 Subject: [PATCH 11/50] Update tests/test_litellm/proxy/management_endpoints/test_ui_sso.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 016131f24dc..097d5f34702 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3205,7 +3205,7 @@ class TestPKCEFunctionality: stored_key = "pkce_verifier:multi_pod_state_xyz" assert stored_key in mock_redis._store stored_value = mock_redis._store[stored_key] - assert isinstance(stored_value, str) and len(stored_value) == 43 + assert isinstance(stored_value, str) and len(json.loads(stored_value)) == 43 # Pod B: callback with same state, retrieve from "Redis" mock_request = MagicMock(spec=Request) From f33a7ca7d6db05a0860dda17a352704ec88806e0 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Sat, 14 Feb 2026 04:34:58 +0530 Subject: [PATCH 12/50] fix: as per request changes --- .../proxy/common_utils/key_rotation_manager.py | 4 ++-- litellm/proxy/utils.py | 18 ++++++++++++++++-- .../common_utils/test_key_rotation_manager.py | 6 +++--- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 88fd08f6ab4..5a0a1fabc7d 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -112,9 +112,9 @@ class KeyRotationManager: result = await self.prisma_client.db.litellm_deprecatedverificationtoken.delete_many( where={"revoke_at": {"lt": now}} ) - if result.get("count", 0) > 0: + if result > 0: verbose_proxy_logger.debug( - "Cleaned up %s expired deprecated key(s)", result["count"] + "Cleaned up %s expired deprecated key(s)", result ) except Exception as e: verbose_proxy_logger.debug( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 528415c891d..4084284f7d0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -76,6 +76,7 @@ from litellm import ( from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache +from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger @@ -2033,7 +2034,8 @@ def jsonify_object(data: dict) -> dict: # In-memory cache for deprecated key lookups: maps old_token_hash -> (active_token_id, expires_at_ts) # Avoids a DB query on every auth request for non-deprecated keys. -_deprecated_key_cache: Dict[str, tuple] = {} +# Bounded to prevent memory leaks from accumulated rotations. +_deprecated_key_cache: LimitedSizeOrderedDict = LimitedSizeOrderedDict(max_size=1000) _DEPRECATED_KEY_CACHE_TTL_SECONDS = 60 @@ -2073,6 +2075,11 @@ async def _lookup_deprecated_key( now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, ) return deprecated_row.active_token_id + # Cache negative result to avoid repeated DB lookups for invalid tokens + _deprecated_key_cache[hashed_token] = ( + None, + now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, + ) except Exception as e: verbose_proxy_logger.debug("Deprecated key lookup skipped: %s", e) @@ -2414,6 +2421,7 @@ class PrismaClient: parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, budget_id_list: Optional[List[str]] = None, + check_deprecated: bool = True, ): args_passed_in = locals() start_time = time.time() @@ -2712,7 +2720,12 @@ class PrismaClient: ) # If not found in main table, check deprecated keys (grace period) - if response is None and hashed_token is not None: + # check_deprecated=False on the recursive call prevents unbounded chaining + if ( + response is None + and hashed_token is not None + and check_deprecated + ): active_token_id = await _lookup_deprecated_key( db=self.db, hashed_token=hashed_token ) @@ -2723,6 +2736,7 @@ class PrismaClient: query_type="find_unique", parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, + check_deprecated=False, ) if response is not None: verbose_proxy_logger.debug( diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py index 002b015d960..24828cdff36 100644 --- a/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_manager.py @@ -233,9 +233,9 @@ class TestKeyRotationManager: Test that _cleanup_expired_deprecated_keys deletes expired deprecated keys. """ mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.return_value = { - "count": 3 - } + mock_prisma_client.db.litellm_deprecatedverificationtoken.delete_many.return_value = ( + 3 + ) manager = KeyRotationManager(mock_prisma_client) await manager._cleanup_expired_deprecated_keys() From d90f3d558ed7fd38df4912309b7803ca6641f475 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sat, 14 Feb 2026 08:46:07 +0530 Subject: [PATCH 13/50] Update litellm/proxy/utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4084284f7d0..f7e32869a23 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2075,11 +2075,8 @@ async def _lookup_deprecated_key( now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, ) return deprecated_row.active_token_id - # Cache negative result to avoid repeated DB lookups for invalid tokens - _deprecated_key_cache[hashed_token] = ( - None, - now_ts + _DEPRECATED_KEY_CACHE_TTL_SECONDS, - ) + # Only cache positive results; negative lookups are fast on indexed columns + # and caching them risks evicting real deprecated key entries. except Exception as e: verbose_proxy_logger.debug("Deprecated key lookup skipped: %s", e) From bac6d1127c15f6a1246ee1536a62140c32d41697 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Fri, 13 Feb 2026 22:18:19 -0500 Subject: [PATCH 14/50] Fix errors when callbacks are invoked for file delete operations: --- litellm/integrations/s3_v2.py | 6 ++- .../proxy/hooks/proxy_track_cost_callback.py | 6 +++ tests/test_litellm/integrations/test_s3_v2.py | 45 +++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 41 +++++++++++++++++ 4 files changed, 97 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 534b85e4752..e847ce27476 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -248,7 +248,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): ) if s3_batch_logging_element is None: - raise ValueError("s3_batch_logging_element is None") + verbose_logger.debug( + "s3 Logging - skipping event, no standard_logging_object for call_type=%s", + kwargs.get("call_type", "unknown"), + ) + return verbose_logger.debug( "\ns3 Logger - Logging payload = %s", s3_batch_logging_element diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 37b79e6d065..ffc9dbd2372 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -202,6 +202,12 @@ class _ProxyDBLogger(CustomLogger): max_budget=end_user_max_budget, ) else: + if sl_object is None and kwargs.get("model") is None: + verbose_proxy_logger.warning( + "Cost tracking - skipping, no standard_logging_object and no model for call_type=%s", + kwargs.get("call_type", "unknown"), + ) + return if kwargs.get("stream") is not True or ( kwargs.get("stream") is True and "complete_streaming_response" in kwargs ): diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 0a3523699a9..51fec288f6e 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -157,6 +157,51 @@ class TestS3V2UnitTests: assert result == {"downloaded": "data"} +@pytest.mark.asyncio +async def test_async_log_event_skips_when_standard_logging_object_missing(): + """ + Reproduces the bug where _async_log_event_base raises ValueError when + kwargs has no standard_logging_object (e.g. call_type=afile_delete). + + The S3 logger should skip gracefully, not raise. + """ + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_region_name="us-east-1", + s3_aws_access_key_id="fake", + s3_aws_secret_access_key="fake", + ) + + kwargs_without_slo = { + "call_type": "afile_delete", + "model": None, + "litellm_call_id": "test-call-id", + } + + start_time = datetime.utcnow() + end_time = datetime.utcnow() + + # Spy on handle_callback_failure — should NOT be called if we skip gracefully. + # Without the fix, the ValueError is caught by the except block which calls + # handle_callback_failure. With the fix, we return early and never hit except. + with patch.object(logger, "handle_callback_failure") as mock_failure: + await logger._async_log_event_base( + kwargs=kwargs_without_slo, + response_obj=None, + start_time=start_time, + end_time=end_time, + ) + + assert not mock_failure.called, ( + "handle_callback_failure should not be called — " + "missing standard_logging_object should be a graceful skip, not an error" + ) + + # Nothing should have been queued (catches the case where code falls + # through without returning and appends None to the queue) + assert len(logger.log_queue) == 0, "log_queue should be empty when standard_logging_object is missing" + + @pytest.mark.asyncio async def test_strip_base64_removes_file_and_nontext_entries(): logger = S3Logger(s3_strip_base64_files=True) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index cb6d90103f7..46da6a89cb2 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -126,3 +126,44 @@ async def test_async_post_call_failure_hook_non_llm_route(): # Assert that update_database was NOT called for non-LLM routes mock_update_database.assert_not_called() + + +@pytest.mark.asyncio +async def test_track_cost_callback_skips_when_no_standard_logging_object(): + """ + Reproduces the bug where _PROXY_track_cost_callback raises + 'Cost tracking failed for model=None' when kwargs has no + standard_logging_object (e.g. call_type=afile_delete). + + File operations have no model and no standard_logging_object. + The callback should skip gracefully instead of raising. + """ + logger = _ProxyDBLogger() + + kwargs = { + "call_type": "afile_delete", + "model": None, + "litellm_call_id": "test-call-id", + "litellm_params": {}, + "stream": False, + } + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + # update_database should NOT be called — nothing to track + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() + + # failed_tracking_alert should NOT be called — this is not an error + mock_proxy_logging.failed_tracking_alert.assert_not_called() From 1b6a7ed6c1b20ce6119d2b7978b6a388dae23f01 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Fri, 13 Feb 2026 23:30:22 -0500 Subject: [PATCH 15/50] Fix errors when callbacks are invoked for file operations --- litellm/_service_logger.py | 2 +- .../proxy/hooks/proxy_track_cost_callback.py | 2 +- .../hooks/test_proxy_track_cost_callback.py | 33 +++++++ tests/test_litellm/test_service_logger.py | 97 +++++++++++++++++++ 4 files changed, 132 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/test_service_logger.py diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index b67d0d86063..e2759e4fab3 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -315,7 +315,7 @@ class ServiceLogging(CustomLogger): await self.async_service_success_hook( service=ServiceTypes.LITELLM, duration=_duration, - call_type=kwargs["call_type"], + call_type=kwargs.get("call_type", "unknown") ) except Exception as e: raise e diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ffc9dbd2372..94886293a08 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -202,7 +202,7 @@ class _ProxyDBLogger(CustomLogger): max_budget=end_user_max_budget, ) else: - if sl_object is None and kwargs.get("model") is None: + if sl_object is None and not kwargs.get("model"): verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object and no model for call_type=%s", kwargs.get("call_type", "unknown"), diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 46da6a89cb2..e8765cf78ca 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -167,3 +167,36 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): # failed_tracking_alert should NOT be called — this is not an error mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model_value", [None, ""]) +async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value): + """ + Same bug as above but model can also be empty string (e.g. health check callbacks). + The guard should catch all falsy model values when sl_object is missing. + """ + logger = _ProxyDBLogger() + + kwargs = { + "call_type": "acompletion", + "model": model_value, + "litellm_params": {}, + "stream": False, + } + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.failed_tracking_alert.assert_not_called() diff --git a/tests/test_litellm/test_service_logger.py b/tests/test_litellm/test_service_logger.py new file mode 100644 index 00000000000..ed44fe9b9f2 --- /dev/null +++ b/tests/test_litellm/test_service_logger.py @@ -0,0 +1,97 @@ +""" +Tests for litellm/_service_logger.py + +Regression test for KeyError: 'call_type' when async_log_success_event +is called without call_type in kwargs (e.g. from batch polling callbacks). +""" + +import pytest +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, patch + +from litellm._service_logger import ServiceLogging + + +@pytest.mark.asyncio +async def test_async_log_success_event_should_not_raise_when_call_type_missing(): + """ + When async_log_success_event is called with kwargs that omit 'call_type', + it should not raise a KeyError. This happens in the batch polling flow + where check_batch_cost.py creates a Logging object whose model_call_details + don't include call_type. + """ + service_logger = ServiceLogging(mock_testing=True) + + start_time = datetime(2026, 2, 13, 22, 35, 0) + end_time = datetime(2026, 2, 13, 22, 35, 1) + kwargs_without_call_type = {"model": "gpt-4", "stream": False} + + with patch.object( + service_logger, "async_service_success_hook", new_callable=AsyncMock + ) as mock_hook: + await service_logger.async_log_success_event( + kwargs=kwargs_without_call_type, + response_obj=None, + start_time=start_time, + end_time=end_time, + ) + + mock_hook.assert_called_once() + call_kwargs = mock_hook.call_args + assert call_kwargs.kwargs["call_type"] == "unknown" + + +@pytest.mark.asyncio +async def test_async_log_success_event_should_pass_call_type_when_present(): + """ + When call_type IS present in kwargs, it should be forwarded correctly. + """ + service_logger = ServiceLogging(mock_testing=True) + + start_time = datetime(2026, 2, 13, 22, 35, 0) + end_time = datetime(2026, 2, 13, 22, 35, 1) + kwargs_with_call_type = { + "model": "gpt-4", + "stream": False, + "call_type": "aretrieve_batch", + } + + with patch.object( + service_logger, "async_service_success_hook", new_callable=AsyncMock + ) as mock_hook: + await service_logger.async_log_success_event( + kwargs=kwargs_with_call_type, + response_obj=None, + start_time=start_time, + end_time=end_time, + ) + + mock_hook.assert_called_once() + call_kwargs = mock_hook.call_args + assert call_kwargs.kwargs["call_type"] == "aretrieve_batch" + + +@pytest.mark.asyncio +async def test_async_log_success_event_should_handle_float_duration(): + """ + When start_time and end_time produce a float duration (not timedelta), + it should still work correctly. + """ + service_logger = ServiceLogging(mock_testing=True) + + start_time = 1000.0 + end_time = 1001.5 + + with patch.object( + service_logger, "async_service_success_hook", new_callable=AsyncMock + ) as mock_hook: + await service_logger.async_log_success_event( + kwargs={"call_type": "completion"}, + response_obj=None, + start_time=start_time, + end_time=end_time, + ) + + mock_hook.assert_called_once() + call_kwargs = mock_hook.call_args + assert call_kwargs.kwargs["duration"] == 1.5 From 358180eb2d96a0116e55fc115089c5fc182448d3 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 00:16:34 -0500 Subject: [PATCH 16/50] Fix: pass deployment credentials to afile_retrieve in managed_files post-call hook --- .../proxy/hooks/managed_files.py | 18 +- .../proxy/test_managed_files_hook.py | 167 ++++++++++++++++++ 2 files changed, 179 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/enterprise/proxy/test_managed_files_hook.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index a41b3f3bf6f..e11e827d454 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -914,12 +914,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Fetch the actual file object from the provider file_object = None try: - # Use litellm to retrieve the file object from the provider - from litellm import afile_retrieve - file_object = await afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", - file_id=original_file_id - ) + from litellm.proxy.proxy_server import llm_router as _llm_router + if _llm_router is not None and model_id: + _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {} + file_object = await litellm.afile_retrieve( + file_id=original_file_id, + **_creds, + ) + else: + file_object = await litellm.afile_retrieve( + custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", + file_id=original_file_id, + ) verbose_logger.debug( f"Successfully retrieved file object for {file_attr}={original_file_id}" ) diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py new file mode 100644 index 00000000000..9526304aff0 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -0,0 +1,167 @@ +""" +Tests for enterprise/litellm_enterprise/proxy/hooks/managed_files.py + +Regression test for afile_retrieve called without credentials in +async_post_call_success_hook when processing completed batch responses. +""" + +import pytest +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.llms.openai import OpenAIFileObject +from litellm.types.utils import LiteLLMBatch + + +def _make_file_object(file_id: str = "file-output-abc") -> OpenAIFileObject: + return OpenAIFileObject( + id=file_id, + bytes=100, + created_at=1700000000, + filename="output.jsonl", + object="file", + purpose="batch_output", + status="processed", + ) + + +def _make_batch_response( + batch_id: str = "batch-123", + output_file_id: Optional[str] = "file-output-abc", + status: str = "completed", + model_id: str = "model-deploy-xyz", + model_name: str = "azure/gpt-4", +) -> LiteLLMBatch: + """Create a LiteLLMBatch response with hidden params set as the router would.""" + batch = LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1700000000, + endpoint="/v1/chat/completions", + input_file_id="file-input-abc", + object="batch", + status=status, + output_file_id=output_file_id, + ) + batch._hidden_params = { + "unified_file_id": "some-unified-id", + "unified_batch_id": "some-unified-batch-id", + "model_id": model_id, + "model_name": model_name, + } + return batch + + +def _make_user_api_key_dict() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id="test-user", + parent_otel_span=None, + ) + + +def _make_managed_files_instance(): + """Create a _PROXY_LiteLLMManagedFiles with storage methods mocked out.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + mock_cache = MagicMock() + mock_prisma = MagicMock() + + instance = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=mock_cache, + prisma_client=mock_prisma, + ) + instance.store_unified_file_id = AsyncMock() + instance.store_unified_object_id = AsyncMock() + return instance + + +@pytest.mark.asyncio +async def test_should_pass_credentials_to_afile_retrieve(): + """ + When async_post_call_success_hook processes a completed batch with an output_file_id, + it calls afile_retrieve to fetch file metadata. It must pass credentials from the + router deployment, not just custom_llm_provider and file_id. + + Regression test for: managed_files.py:919 calling afile_retrieve without api_key/api_base. + """ + managed_files = _make_managed_files_instance() + batch_response = _make_batch_response( + model_id="model-deploy-xyz", + model_name="azure/gpt-4", + output_file_id="file-output-abc", + ) + user_api_key_dict = _make_user_api_key_dict() + + mock_credentials = { + "api_key": "test-azure-key", + "api_base": "https://my-azure.openai.azure.com/", + "api_version": "2025-03-01-preview", + "custom_llm_provider": "azure", + } + + mock_router = MagicMock() + mock_router.get_deployment_credentials_with_provider = MagicMock( + return_value=mock_credentials + ) + + mock_afile_retrieve = AsyncMock(return_value=_make_file_object("file-output-abc")) + + with patch( + "litellm.afile_retrieve", mock_afile_retrieve + ), patch( + "litellm.proxy.proxy_server.llm_router", mock_router + ): + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response=batch_response, + ) + + mock_afile_retrieve.assert_called() + call_kwargs = mock_afile_retrieve.call_args + + assert call_kwargs.kwargs.get("api_key") == "test-azure-key", ( + f"afile_retrieve must receive api_key from router credentials. " + f"Got kwargs: {call_kwargs.kwargs}" + ) + assert call_kwargs.kwargs.get("api_base") == "https://my-azure.openai.azure.com/", ( + f"afile_retrieve must receive api_base from router credentials. " + f"Got kwargs: {call_kwargs.kwargs}" + ) + + +@pytest.mark.asyncio +async def test_should_fallback_when_no_router(): + """ + When llm_router is not available, afile_retrieve should still be called + with the fallback behavior (custom_llm_provider extracted from model_name). + """ + managed_files = _make_managed_files_instance() + batch_response = _make_batch_response( + model_id="model-deploy-xyz", + model_name="azure/gpt-4", + output_file_id="file-output-abc", + ) + user_api_key_dict = _make_user_api_key_dict() + + mock_afile_retrieve = AsyncMock(return_value=_make_file_object("file-output-abc")) + + with patch( + "litellm.afile_retrieve", mock_afile_retrieve + ), patch( + "litellm.proxy.proxy_server.llm_router", None + ): + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response=batch_response, + ) + + mock_afile_retrieve.assert_called() + call_kwargs = mock_afile_retrieve.call_args + assert call_kwargs.kwargs.get("custom_llm_provider") == "azure" + assert call_kwargs.kwargs.get("file_id") == "file-output-abc" From 5433ae7e8caf725191bca6a579596b404f4acf72 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 00:30:35 -0500 Subject: [PATCH 17/50] Fix: bypass managed files access check in batch polling by calling afile_content directly --- .../proxy/common_utils/check_batch_cost.py | 39 ++-- .../proxy/test_managed_files_access_check.py | 200 ++++++++++++++++++ 2 files changed, 220 insertions(+), 19 deletions(-) create mode 100644 tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index bb25e4f0626..807d237df5c 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -4,7 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t from litellm._uuid import uuid from datetime import datetime -from typing import TYPE_CHECKING, Optional, cast +from typing import TYPE_CHECKING, Optional from litellm._logging import verbose_proxy_logger @@ -35,14 +35,11 @@ class CheckBatchCost: - if not, return False - if so, return True """ - from litellm_enterprise.proxy.hooks.managed_files import ( - _PROXY_LiteLLMManagedFiles, - ) - from litellm.batches.batch_utils import ( _get_file_content_as_dictionary, calculate_batch_cost_and_usage, ) + from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -102,27 +99,31 @@ class CheckBatchCost: continue ## RETRIEVE THE BATCH JOB OUTPUT FILE - managed_files_obj = cast( - Optional[_PROXY_LiteLLMManagedFiles], - self.proxy_logging_obj.get_proxy_hook("managed_files"), - ) if ( response.status == "completed" and response.output_file_id is not None - and managed_files_obj is not None ): verbose_proxy_logger.info( f"Batch ID: {batch_id} is complete, tracking cost and usage" ) - # track cost - model_file_id_mapping = { - response.output_file_id: {model_id: response.output_file_id} - } - _file_content = await managed_files_obj.afile_content( - file_id=response.output_file_id, - litellm_parent_otel_span=None, - llm_router=self.llm_router, - model_file_id_mapping=model_file_id_mapping, + + # Extract raw provider file ID from the unified file ID + # (async_post_call_success_hook may have replaced output_file_id with a unified ID) + raw_output_file_id = response.output_file_id + decoded = _is_base64_encoded_unified_file_id(raw_output_file_id) + if decoded: + try: + raw_output_file_id = decoded.split("llm_output_file_id,")[1].split(";")[0] + except (IndexError, AttributeError): + pass + + # Call litellm.afile_content directly with deployment credentials, + # bypassing the managed files access-control hooks that would + # reject this background job's default_user_id identity + credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} + _file_content = await afile_content( + file_id=raw_output_file_id, + **credentials, ) file_content_as_dict = _get_file_content_as_dictionary( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py new file mode 100644 index 00000000000..2db5a2214cb --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -0,0 +1,200 @@ +""" +Tests for managed files access control in batch polling context. + +Regression test for: batch polling job running as default_user_id gets 403 +when trying to access managed files created by a real user. + +The fix (Option C) makes check_batch_cost call litellm.afile_content directly +with deployment credentials, bypassing the managed files access-control hooks. +""" + +import base64 +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth + + +def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + parent_otel_span=None, + ) + + +def _make_unified_file_id() -> str: + """Create a base64-encoded unified file ID that passes _is_base64_encoded_unified_file_id.""" + raw = "litellm_proxy:application/octet-stream;unified_id,test-123;target_model_names,azure-gpt-4" + return base64.b64encode(raw.encode()).decode() + + +def _make_managed_files_instance(file_created_by: str, unified_file_id: str): + """Create a _PROXY_LiteLLMManagedFiles with a mocked DB that returns a file owned by file_created_by.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + mock_db_record = MagicMock() + mock_db_record.created_by = file_created_by + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=mock_db_record + ) + + instance = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=mock_prisma, + ) + return instance + + +# --- Access control unit tests (document existing behavior) --- + + +@pytest.mark.asyncio +async def test_should_allow_file_owner_access(): + """File owner can access their own file — baseline sanity check.""" + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_instance( + file_created_by="user-A", + unified_file_id=unified_file_id, + ) + user = _make_user_api_key_dict("user-A") + data = {"file_id": unified_file_id} + + result = await managed_files.check_managed_file_id_access(data, user) + assert result is True + + +@pytest.mark.asyncio +async def test_should_block_different_user_access(): + """A different regular user cannot access another user's file — correct behavior.""" + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_instance( + file_created_by="user-A", + unified_file_id=unified_file_id, + ) + user = _make_user_api_key_dict("user-B") + data = {"file_id": unified_file_id} + + with pytest.raises(HTTPException) as exc_info: + await managed_files.check_managed_file_id_access(data, user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_should_block_default_user_id_access(): + """ + default_user_id is correctly blocked by the access check. + This documents the existing behavior that the Option C fix works around. + """ + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_instance( + file_created_by="user-A", + unified_file_id=unified_file_id, + ) + system_user = _make_user_api_key_dict("default_user_id") + data = {"file_id": unified_file_id} + + with pytest.raises(HTTPException) as exc_info: + await managed_files.check_managed_file_id_access(data, system_user) + assert exc_info.value.status_code == 403 + + +# --- Option C fix test: check_batch_cost bypasses managed files hook --- + + +@pytest.mark.asyncio +async def test_check_batch_cost_should_call_afile_content_directly_with_credentials(): + """ + check_batch_cost should call litellm.afile_content directly with deployment + credentials, bypassing managed_files_obj.afile_content and its access-control + hooks. This avoids the 403 that occurs when the background job runs as + default_user_id. + """ + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + # Build a unified object ID in the expected format: + # litellm_proxy;model_id:{};llm_batch_id:{};llm_output_file_id:{} + unified_raw = "litellm_proxy;model_id:model-deploy-xyz;llm_batch_id:batch-123;llm_output_file_id:file-raw-output" + unified_object_id = base64.b64encode(unified_raw.encode()).decode() + + # Mock a pending job from the DB + mock_job = MagicMock() + mock_job.unified_object_id = unified_object_id + mock_job.created_by = "user-A" + mock_job.id = "job-1" + + # Mock prisma + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + + # Mock proxy_logging_obj — should NOT be called for file content + mock_proxy_logging = MagicMock() + mock_managed_files_hook = MagicMock() + mock_managed_files_hook.afile_content = AsyncMock() + mock_proxy_logging.get_proxy_hook = MagicMock(return_value=mock_managed_files_hook) + + # Mock the batch response (completed, with output file) + from litellm.types.utils import LiteLLMBatch + batch_response = LiteLLMBatch( + id="batch-123", + completion_window="24h", + created_at=1700000000, + endpoint="/v1/chat/completions", + input_file_id="file-input", + object="batch", + status="completed", + output_file_id="file-raw-output", + ) + + # Mock router + mock_router = MagicMock() + mock_router.aretrieve_batch = AsyncMock(return_value=batch_response) + mock_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "api_key": "test-key", + "api_base": "https://test.azure.com/", + "custom_llm_provider": "azure", + } + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "azure" + mock_deployment.litellm_params.model = "azure/gpt-4" + mock_router.get_deployment = MagicMock(return_value=mock_deployment) + + checker = CheckBatchCost( + proxy_logging_obj=mock_proxy_logging, + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"id":"req-1","response":{"status_code":200,"body":{"id":"cmpl-1","object":"chat.completion","created":1700000000,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}}}\n' + + with patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ) as mock_direct_afile_content: + await checker.check_batch_cost() + + # afile_content should be called directly (not through managed_files_obj) + mock_direct_afile_content.assert_called_once() + call_kwargs = mock_direct_afile_content.call_args.kwargs + + assert call_kwargs.get("api_key") == "test-key", ( + f"afile_content should receive api_key from deployment credentials. " + f"Got: {call_kwargs}" + ) + + # managed_files_obj.afile_content should NOT have been called + mock_managed_files_hook.afile_content.assert_not_called() From 3743513bb64bfcd5fe0321953f073c06565d5620 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sat, 14 Feb 2026 11:01:25 +0530 Subject: [PATCH 18/50] Update tests/test_litellm/proxy/management_endpoints/test_ui_sso.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/management_endpoints/test_ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 097d5f34702..938d4e8c871 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3079,7 +3079,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.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): # Act token_params = ( await SSOAuthenticationHandler.prepare_token_exchange_parameters( From e59c8d22afe9d1c2ec67fd4f070346e06c3461f2 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 10:07:12 -0500 Subject: [PATCH 19/50] fix: afile_retrieve returns unified ID for batch output files --- .../proxy/hooks/managed_files.py | 1 + fix_afile_retrieve_returns_unified_id.md | 47 +++++++++++++ .../test_afile_retrieve_returns_unified_id.py | 67 +++++++++++++++++++ 3 files changed, 115 insertions(+) create mode 100644 fix_afile_retrieve_returns_unified_id.md create mode 100644 tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index e11e827d454..32d7313aae4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1011,6 +1011,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Case 2: Managed file and the file object exists in the database if stored_file_object and stored_file_object.file_object: + stored_file_object.file_object.id = file_id return stored_file_object.file_object # Case 3: Managed file exists in the database but not the file object (for. e.g the batch task might not have run) diff --git a/fix_afile_retrieve_returns_unified_id.md b/fix_afile_retrieve_returns_unified_id.md new file mode 100644 index 00000000000..c8cc72d6f99 --- /dev/null +++ b/fix_afile_retrieve_returns_unified_id.md @@ -0,0 +1,47 @@ +# Fix: afile_retrieve returns raw provider ID for batch output files + +## Bug + +`managed_files.afile_retrieve()` Case 2 (file_object already in DB) returned the stored `file_object` without replacing `.id` with the unified file ID. Case 3 (fetch from provider) did this correctly at line 1028. + +## Fix + +One-line change in `enterprise/litellm_enterprise/proxy/hooks/managed_files.py`: + +```python +# Before (line 1013-1014) +if stored_file_object and stored_file_object.file_object: + return stored_file_object.file_object + +# After +if stored_file_object and stored_file_object.file_object: + stored_file_object.file_object.id = file_id + return stored_file_object.file_object +``` + +## Test + +```bash +poetry run pytest tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py -s -vvvv +``` + +## Test failure (before fix) + +``` +FAILED tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py::test_should_return_unified_id_when_file_object_exists_in_db +AssertionError: afile_retrieve should return the unified ID 'bGl0ZWxsbV9wcm94eTp1bmlmaWVkX291dHB1dF9maWxl', but got raw provider ID 'batch_20260214-output-file-1' +assert 'batch_20260214-output-file-1' == 'bGl0ZWxsbV9wcm94eTp1bmlmaWVkX291dHB1dF9maWxl' +=================== 1 failed, 1 retried in 102.95s =================== +``` + +## Test pass (after fix) + +``` +tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py::test_should_return_unified_id_when_file_object_exists_in_db PASSED +============================== 1 passed in 0.11s =============================== +``` + +## Files changed + +- `enterprise/litellm_enterprise/proxy/hooks/managed_files.py` — one-line fix +- `tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py` — new test diff --git a/tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py b/tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py new file mode 100644 index 00000000000..7040aef73e5 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py @@ -0,0 +1,67 @@ +""" +Test that managed_files.afile_retrieve returns the unified file ID, not the +raw provider file ID, when file_object is already stored in the database. + +Bug: managed_files.py Case 2 returns stored_file_object.file_object directly +without replacing .id with the unified ID. Case 3 (fetch from provider) does +it correctly at line 1028. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from litellm.proxy._types import LiteLLM_ManagedFileTable +from litellm.types.llms.openai import OpenAIFileObject + + +def _make_managed_files_instance(): + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + instance = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=MagicMock(), + ) + return instance + + +@pytest.mark.asyncio +async def test_should_return_unified_id_when_file_object_exists_in_db(): + """ + When get_unified_file_id returns a stored file_object (Case 2), + afile_retrieve must set .id to the unified file ID before returning. + """ + unified_id = "bGl0ZWxsbV9wcm94eTp1bmlmaWVkX291dHB1dF9maWxl" + raw_provider_id = "batch_20260214-output-file-1" + + stored = LiteLLM_ManagedFileTable( + unified_file_id=unified_id, + file_object=OpenAIFileObject( + id=raw_provider_id, + bytes=489, + created_at=1700000000, + filename="batch_output.jsonl", + object="file", + purpose="batch_output", + status="processed", + ), + model_mappings={"model-abc": raw_provider_id}, + flat_model_file_ids=[raw_provider_id], + created_by="test-user", + updated_by="test-user", + ) + + managed_files = _make_managed_files_instance() + managed_files.get_unified_file_id = AsyncMock(return_value=stored) + + result = await managed_files.afile_retrieve( + file_id=unified_id, + litellm_parent_otel_span=None, + llm_router=None, + ) + + assert result.id == unified_id, ( + f"afile_retrieve should return the unified ID '{unified_id}', " + f"but got raw provider ID '{result.id}'" + ) From cd0fed826bfd22b197dbfdc0763e0eaa3f43e5cc Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 11:24:08 -0500 Subject: [PATCH 20/50] fix: batch retrieve returns unified input_file_id --- litellm/proxy/batches_endpoints/endpoints.py | 34 +++++ .../openai_files_endpoints/common_utils.py | 10 ++ .../test_batch_retrieve_input_file_id.py | 75 +++++++++++ ..._retrieve_returns_unified_input_file_id.py | 124 ++++++++++++++++++ 4 files changed, 243 insertions(+) create mode 100644 tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py create mode 100644 tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 06800cb4524..25680ab1220 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -377,6 +377,23 @@ async def retrieve_batch( response = await proxy_logging_obj.post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response ) + + # Resolve raw input_file_id to unified ID + if ( + unified_batch_id + and hasattr(response, "input_file_id") + and response.input_file_id + and not _is_base64_encoded_unified_file_id(response.input_file_id) + and prisma_client + ): + try: + _managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if _managed_file: + response.input_file_id = _managed_file.unified_file_id + except Exception: + pass asyncio.create_task( proxy_logging_obj.update_request_status( @@ -479,6 +496,23 @@ async def retrieve_batch( data=data, user_api_key_dict=user_api_key_dict, response=response ) + # Resolve raw input_file_id to unified ID + if ( + unified_batch_id + and hasattr(response, "input_file_id") + and response.input_file_id + and not _is_base64_encoded_unified_file_id(response.input_file_id) + and prisma_client + ): + try: + _managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if _managed_file: + response.input_file_id = _managed_file.unified_file_id + except Exception: + pass + ### ALERTING ### asyncio.create_task( proxy_logging_obj.update_request_status( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f67dc5e2aaa..15c5bdfabb2 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -687,6 +687,16 @@ async def get_batch_from_database( batch_data = json.loads(db_batch_object.file_object) if isinstance(db_batch_object.file_object, str) else db_batch_object.file_object response = LiteLLMBatch(**batch_data) response.id = batch_id + + if response.input_file_id and not _is_base64_encoded_unified_file_id(response.input_file_id): + try: + managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if managed_file: + response.input_file_id = managed_file.unified_file_id + except Exception: + pass verbose_proxy_logger.debug( f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}" diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py new file mode 100644 index 00000000000..6e9c3c0354b --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py @@ -0,0 +1,75 @@ +""" +Test that batch retrieve endpoint resolves raw input_file_id to the +unified managed file ID before returning. + +Bug: After batch completion, batches.retrieve returns the raw provider +input_file_id instead of the LiteLLM unified ID. +""" + +import base64 +import json + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, +) + + +DECODED_UNIFIED_INPUT_FILE_ID = "litellm_proxy:application/octet-stream;unified_id,test-uuid;target_model_names,azure-gpt-4" +B64_UNIFIED_INPUT_FILE_ID = base64.urlsafe_b64encode(DECODED_UNIFIED_INPUT_FILE_ID.encode()).decode().rstrip("=") +RAW_INPUT_FILE_ID = "file-raw-provider-abc123" + +DECODED_UNIFIED_BATCH_ID = "litellm_proxy;model_id:model-xyz;llm_batch_id:batch-123" +B64_UNIFIED_BATCH_ID = base64.urlsafe_b64encode(DECODED_UNIFIED_BATCH_ID.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +async def test_should_resolve_raw_input_file_id_to_unified(): + """ + When a completed batch has a raw input_file_id and the managed file table + contains a record for that raw ID, the retrieve endpoint should resolve + it to the unified file ID. + """ + unified_batch_id = _is_base64_encoded_unified_file_id(B64_UNIFIED_BATCH_ID) + assert unified_batch_id, "Test setup: batch_id should decode as unified" + + from litellm.types.utils import LiteLLMBatch + + batch_data = { + "id": B64_UNIFIED_BATCH_ID, + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": RAW_INPUT_FILE_ID, + "object": "batch", + "status": "completed", + "output_file_id": "file-output-xyz", + } + + mock_db_object = MagicMock() + mock_db_object.file_object = json.dumps(batch_data) + + mock_managed_file = MagicMock() + mock_managed_file.unified_file_id = B64_UNIFIED_INPUT_FILE_ID + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=mock_db_object) + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=mock_managed_file) + + from litellm.proxy.openai_files_endpoints.common_utils import get_batch_from_database + + _, response = await get_batch_from_database( + batch_id=B64_UNIFIED_BATCH_ID, + unified_batch_id=unified_batch_id, + managed_files_obj=MagicMock(), + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + ) + + assert response is not None, "Batch should be found in DB" + assert response.input_file_id == B64_UNIFIED_INPUT_FILE_ID, ( + f"input_file_id should be unified '{B64_UNIFIED_INPUT_FILE_ID}', " + f"got raw '{response.input_file_id}'" + ) diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py new file mode 100644 index 00000000000..420f5f9789c --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py @@ -0,0 +1,124 @@ +""" +Test that get_batch_from_database resolves raw input_file_id to the +unified/managed file ID when reading a batch from the database. + +Bug: The batch retrieve path stores the raw provider input_file_id in the +DB (via async_post_call_success_hook on the retrieve endpoint). When the +batch is later read from DB, get_batch_from_database returns the raw ID +without resolving it to the unified ID. +""" + +import json +import pytest +from typing import Optional +from unittest.mock import AsyncMock, MagicMock + +from litellm.proxy.openai_files_endpoints.common_utils import get_batch_from_database + + +def _mock_prisma(batch_json: str, managed_file_record=None): + """Create a mock prisma client with canned responses.""" + prisma = MagicMock() + + batch_db_record = MagicMock() + batch_db_record.file_object = batch_json + + prisma.db.litellm_managedobjecttable.find_first = AsyncMock( + return_value=batch_db_record + ) + + prisma.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=managed_file_record + ) + + return prisma + + +@pytest.mark.asyncio +async def test_should_resolve_raw_input_file_id_to_unified_id(): + """ + When input_file_id in the stored batch is a raw provider ID, + get_batch_from_database must look up the unified ID from the + managed files table. + """ + unified_batch_id = "bGl0ZWxsbV9wcm94eTpiYXRjaF9pZA" + unified_input_file_id = "bGl0ZWxsbV9wcm94eTp1bmlmaWVkX2lucHV0" + raw_input_file_id = "file-abc123-raw" + + batch_data = { + "id": "batch-raw-123", + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": raw_input_file_id, + "object": "batch", + "status": "completed", + "output_file_id": "file-output-raw", + } + + managed_file_record = MagicMock() + managed_file_record.unified_file_id = unified_input_file_id + + prisma = _mock_prisma( + batch_json=json.dumps(batch_data), + managed_file_record=managed_file_record, + ) + + _, response = await get_batch_from_database( + batch_id=unified_batch_id, + unified_batch_id="decoded_unified_batch_id", + managed_files_obj=MagicMock(), + prisma_client=prisma, + verbose_proxy_logger=MagicMock(), + ) + + assert response is not None + assert response.input_file_id == unified_input_file_id, ( + f"input_file_id should be resolved to '{unified_input_file_id}', " + f"got raw: '{response.input_file_id}'" + ) + + prisma.db.litellm_managedfiletable.find_first.assert_called_once_with( + where={"flat_model_file_ids": {"has": raw_input_file_id}} + ) + + +@pytest.mark.asyncio +async def test_should_preserve_already_managed_input_file_id(): + """ + When input_file_id is already a managed/unified ID, it should + not be modified. + """ + import base64 + + unified_batch_id = "bGl0ZWxsbV9wcm94eTpiYXRjaF9pZA" + decoded_unified = "litellm_proxy:application/octet-stream;unified_id,test-123" + base64_input_file_id = base64.urlsafe_b64encode(decoded_unified.encode()).decode().rstrip("=") + + batch_data = { + "id": "batch-raw-123", + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": base64_input_file_id, + "object": "batch", + "status": "completed", + } + + prisma = _mock_prisma(batch_json=json.dumps(batch_data)) + + _, response = await get_batch_from_database( + batch_id=unified_batch_id, + unified_batch_id="decoded_unified_batch_id", + managed_files_obj=MagicMock(), + prisma_client=prisma, + verbose_proxy_logger=MagicMock(), + ) + + assert response is not None + assert response.input_file_id == base64_input_file_id, ( + f"input_file_id was already managed, should be preserved as '{base64_input_file_id}', " + f"got: '{response.input_file_id}'" + ) + + prisma.db.litellm_managedfiletable.find_first.assert_not_called() From ab2d6e4eac89ee562cd3c608641a0b182bb733e1 Mon Sep 17 00:00:00 2001 From: Jay Prajapati <79649559+jayy-77@users.noreply.github.com> Date: Sat, 14 Feb 2026 23:48:49 +0530 Subject: [PATCH 21/50] fix(chatgpt): drop unsupported responses params for Codex Co-authored-by: Cursor --- .../llms/chatgpt/responses/transformation.py | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 0ce24f63a89..bcb6edd39f9 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -73,10 +73,6 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params, headers, ) - request.pop("max_output_tokens", None) - request.pop("max_tokens", None) - request.pop("max_completion_tokens", None) - request.pop("metadata", None) base_instructions = get_chatgpt_default_instructions() existing_instructions = request.get("instructions") if existing_instructions: @@ -92,7 +88,22 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): if "reasoning.encrypted_content" not in include: include.append("reasoning.encrypted_content") request["include"] = include - return request + + allowed_keys = { + "model", + "input", + "instructions", + "stream", + "store", + "include", + "tools", + "tool_choice", + "reasoning", + "previous_response_id", + "truncation", + } + + return {k: v for k, v in request.items() if k in allowed_keys} def transform_response_api_response( self, From 35bbcf9e9889ecf0e41d81f6543f506ee9c055a7 Mon Sep 17 00:00:00 2001 From: Jay Prajapati <79649559+jayy-77@users.noreply.github.com> Date: Sat, 14 Feb 2026 23:48:53 +0530 Subject: [PATCH 22/50] test(chatgpt): ensure Codex request filters unsupported params Co-authored-by: Cursor --- .../test_chatgpt_responses_transformation.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index bec748d8dc8..03cea8785bc 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -88,6 +88,45 @@ class TestChatGPTResponsesAPITransformation: "You are Codex, based on GPT-5." ) + def test_chatgpt_drops_unsupported_responses_params(self): + config = ChatGPTResponsesAPIConfig() + request = config.transform_responses_api_request( + model="chatgpt/gpt-5.2-codex", + input="hi", + response_api_optional_request_params={ + # unsupported by ChatGPT Codex + "user": "user_123", + "temperature": 0.2, + "top_p": 0.9, + "context_management": [{"type": "compaction", "compact_threshold": 200000}], + "metadata": {"foo": "bar"}, + "max_output_tokens": 123, + "stream_options": {"include_usage": True}, + # supported and should be preserved + "truncation": "auto", + "previous_response_id": "resp_123", + "reasoning": {"effort": "medium"}, + "tools": [{"type": "function", "function": {"name": "hello"}}], + "tool_choice": {"type": "function", "function": {"name": "hello"}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "user" not in request + assert "temperature" not in request + assert "top_p" not in request + assert "context_management" not in request + assert "metadata" not in request + assert "max_output_tokens" not in request + assert "stream_options" not in request + + assert request["truncation"] == "auto" + assert request["previous_response_id"] == "resp_123" + assert request["reasoning"] == {"effort": "medium"} + assert request["tools"] == [{"type": "function", "function": {"name": "hello"}}] + assert request["tool_choice"] == {"type": "function", "function": {"name": "hello"}} + def test_chatgpt_non_stream_sse_response_parsing(self): config = ChatGPTResponsesAPIConfig() response_payload = { From 4d87cb8fe306578dda4d1183eaf9c86f035911de Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 19:12:03 -0500 Subject: [PATCH 23/50] Fix deleted managed files returning 403 instead of 404 --- .../proxy/hooks/managed_files.py | 2 +- .../test_deleted_file_returns_403_not_404.py | 124 ++++++++++++++++++ 2 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 32d7313aae4..9601370ef3e 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -230,7 +230,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if managed_file: return managed_file.created_by == user_id - return False + return True # allow through when record not found — downstream will return 404 async def can_user_call_unified_object_id( self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth diff --git a/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py b/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py new file mode 100644 index 00000000000..042796c9640 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py @@ -0,0 +1,124 @@ +""" +Regression test: deleted managed files should not return 403. + +When a managed file's DB record has been deleted, can_user_call_unified_file_id() +returns False (record not found → treated as "access denied"). This causes +check_managed_file_id_access() to raise 403 instead of allowing the request +through so downstream code can return a proper 404. + +The equivalent method for objects (can_user_call_unified_object_id) already +returns True when the record is missing. +""" + +import base64 + +import pytest +from unittest.mock import AsyncMock, MagicMock + +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth + + +def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + parent_otel_span=None, + ) + + +def _make_unified_file_id() -> str: + raw = "litellm_proxy:application/octet-stream;unified_id,test-deleted-file;target_model_names,azure-gpt-4" + return base64.b64encode(raw.encode()).decode() + + +def _make_managed_files_with_no_db_record(): + """Create a _PROXY_LiteLLMManagedFiles where the DB returns None (file was deleted).""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + + return _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=mock_prisma, + ) + + +@pytest.mark.asyncio +async def test_should_not_raise_403_for_deleted_file(): + """ + When a managed file record has been deleted from the DB, + check_managed_file_id_access should NOT raise 403. + It should allow the request through so downstream can return 404. + """ + unified_file_id = _make_unified_file_id() + managed_files = _make_managed_files_with_no_db_record() + user = _make_user_api_key_dict("any-user") + data = {"file_id": unified_file_id} + + # This should NOT raise — deleted file should pass through access check + result = await managed_files.check_managed_file_id_access(data, user) + assert result is True + + +@pytest.mark.asyncio +async def test_should_allow_owner_access_when_record_exists(): + """Baseline: file owner can access their own file.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + unified_file_id = _make_unified_file_id() + + mock_db_record = MagicMock() + mock_db_record.created_by = "user-A" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=mock_db_record + ) + + managed_files = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=mock_prisma, + ) + + user = _make_user_api_key_dict("user-A") + data = {"file_id": unified_file_id} + + result = await managed_files.check_managed_file_id_access(data, user) + assert result is True + + +@pytest.mark.asyncio +async def test_should_block_different_user_when_record_exists(): + """Baseline: different user cannot access another user's file.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + unified_file_id = _make_unified_file_id() + + mock_db_record = MagicMock() + mock_db_record.created_by = "user-A" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=mock_db_record + ) + + managed_files = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=MagicMock(), + prisma_client=mock_prisma, + ) + + user = _make_user_api_key_dict("user-B") + data = {"file_id": unified_file_id} + + with pytest.raises(HTTPException) as exc_info: + await managed_files.check_managed_file_id_access(data, user) + assert exc_info.value.status_code == 403 From a5626768a307289186973252a4fe008281eae869 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 14 Feb 2026 19:46:56 -0500 Subject: [PATCH 24/50] Add comments --- .../proxy/common_utils/check_batch_cost.py | 8 +++----- .../litellm_enterprise/proxy/hooks/managed_files.py | 9 +++++++-- litellm/_service_logger.py | 2 ++ litellm/integrations/s3_v2.py | 2 ++ litellm/proxy/batches_endpoints/endpoints.py | 6 ++++-- litellm/proxy/hooks/batch_rate_limiter.py | 3 ++- litellm/proxy/hooks/proxy_track_cost_callback.py | 2 ++ litellm/proxy/openai_files_endpoints/common_utils.py | 1 + 8 files changed, 23 insertions(+), 10 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 807d237df5c..8e4a154c3cb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -107,8 +107,9 @@ class CheckBatchCost: f"Batch ID: {batch_id} is complete, tracking cost and usage" ) - # Extract raw provider file ID from the unified file ID - # (async_post_call_success_hook may have replaced output_file_id with a unified ID) + # This background job runs as default_user_id, so going through the HTTP endpoint + # would trigger check_managed_file_id_access and get 403. Instead, extract the raw + # provider file ID and call afile_content directly with deployment credentials. raw_output_file_id = response.output_file_id decoded = _is_base64_encoded_unified_file_id(raw_output_file_id) if decoded: @@ -117,9 +118,6 @@ class CheckBatchCost: except (IndexError, AttributeError): pass - # Call litellm.afile_content directly with deployment credentials, - # bypassing the managed files access-control hooks that would - # reject this background job's default_user_id identity credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} _file_content = await afile_content( file_id=raw_output_file_id, diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 9601370ef3e..d9514bf28bc 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -230,7 +230,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if managed_file: return managed_file.created_by == user_id - return True # allow through when record not found — downstream will return 404 + # When DB record is missing (file was deleted), allow through so downstream returns 404. + # Matches can_user_call_unified_object_id which also returns True for missing records. + return True async def can_user_call_unified_object_id( self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth @@ -911,7 +913,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) setattr(response, file_attr, unified_file_id) - # Fetch the actual file object from the provider + # Use llm_router credentials when available. Without credentials, + # Azure and other auth-required providers return 500/401. file_object = None try: from litellm.proxy.proxy_server import llm_router as _llm_router @@ -1010,6 +1013,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): raise Exception(f"LiteLLM Managed File object with id={file_id} not found") # Case 2: Managed file and the file object exists in the database + # The stored file_object has the raw provider ID. Replace with the unified ID + # so callers see a consistent ID (matching Case 3 which does response.id = file_id). if stored_file_object and stored_file_object.file_object: stored_file_object.file_object.id = file_id return stored_file_object.file_object diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index e2759e4fab3..8f9a3c5083f 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -312,6 +312,8 @@ class ServiceLogging(CustomLogger): _duration, type(_duration) ) ) # invalid _duration value + # Batch polling callbacks (check_batch_cost) don't include call_type in kwargs. + # Use .get() to avoid KeyError. await self.async_service_success_hook( service=ServiceTypes.LITELLM, duration=_duration, diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index e847ce27476..e0932fc3373 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -247,6 +247,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): standard_logging_payload=kwargs.get("standard_logging_object", None), ) + # afile_delete and other non-model call types never produce a standard_logging_object, + # so s3_batch_logging_element is None. Skip gracefully instead of raising ValueError. if s3_batch_logging_element is None: verbose_logger.debug( "s3 Logging - skipping event, no standard_logging_object for call_type=%s", diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 25680ab1220..b125b90338b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -378,7 +378,8 @@ async def retrieve_batch( data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Resolve raw input_file_id to unified ID + # async_post_call_success_hook replaces batch.id and output_file_id with unified IDs + # but not input_file_id. Look up the unified ID from flat_model_file_ids. if ( unified_batch_id and hasattr(response, "input_file_id") @@ -496,7 +497,8 @@ async def retrieve_batch( data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Resolve raw input_file_id to unified ID + # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id (terminal state path) + # Same as above — resolve raw provider input_file_id to unified ID. if ( unified_batch_id and hasattr(response, "input_file_id") diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 45b1bd8653f..5bebcc92072 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -259,9 +259,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) + # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) + # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) if is_managed_file and user_api_key_dict is not None: - # For managed files, use the managed files hook directly file_content = await self._fetch_managed_file_content( file_id=file_id, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 94886293a08..d903ce0d9d7 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -202,6 +202,8 @@ class _ProxyDBLogger(CustomLogger): max_budget=end_user_max_budget, ) else: + # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. + # Use .get() for "stream" to avoid KeyError on health checks. if sl_object is None and not kwargs.get("model"): verbose_proxy_logger.warning( "Cost tracking - skipping, no standard_logging_object and no model for call_type=%s", diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 15c5bdfabb2..0e3b31b2aa7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -688,6 +688,7 @@ async def get_batch_from_database( response = LiteLLMBatch(**batch_data) response.id = batch_id + # The stored batch object has the raw provider input_file_id. Resolve to unified ID. if response.input_file_id and not _is_base64_encoded_unified_file_id(response.input_file_id): try: managed_file = await prisma_client.db.litellm_managedfiletable.find_first( From 2c654bd7dd34ca922ba255d7f372073ecbcbbb38 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 15 Feb 2026 08:05:44 +0530 Subject: [PATCH 25/50] Update litellm/proxy/utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f7e32869a23..bb95b98da66 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2054,9 +2054,10 @@ async def _lookup_deprecated_key( # Check cache first cached = _deprecated_key_cache.get(hashed_token) + cached = _deprecated_key_cache.get(hashed_token) if cached is not None: - active_token_id, cache_expires_at_ts = cached - if now_ts < cache_expires_at_ts: + active_token_id, cache_expires_at_ts, revoke_at_ts = cached + if now_ts < cache_expires_at_ts and now_ts < revoke_at_ts: return active_token_id else: _deprecated_key_cache.pop(hashed_token, None) From 7d794b567c31a866bc744ec2efd4f7ac10b9fb19 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sun, 15 Feb 2026 14:53:30 -0500 Subject: [PATCH 26/50] fix: thread deployment model_info through batch cost calculation batch_cost_calculator only checked the global cost map, ignoring deployment-level custom pricing (input_cost_per_token_batches etc.). Add optional model_info param through the batch cost chain and pass it from CheckBatchCost. --- .../proxy/common_utils/check_batch_cost.py | 4 + litellm/batches/batch_utils.py | 36 ++++- litellm/cost_calculator.py | 22 ++- .../test_batch_custom_pricing.py | 131 ++++++++++++++++++ 4 files changed, 180 insertions(+), 13 deletions(-) create mode 100644 tests/batches_tests/test_batch_custom_pricing.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 8e4a154c3cb..b28b4497e7c 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -142,11 +142,15 @@ class CheckBatchCost: custom_llm_provider=custom_llm_provider, ) + # Pass deployment model_info so custom batch pricing + # (input_cost_per_token_batches etc.) is used for cost calc + deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {} batch_cost, batch_usage, batch_models = ( await calculate_batch_cost_and_usage( file_content_dictionary=file_content_as_dict, custom_llm_provider=llm_provider, # type: ignore model_name=model_name, + model_info=deployment_model_info, ) ) logging_obj = LiteLLMLogging( diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 16a467e00cb..c92ab9b230e 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -16,14 +16,22 @@ async def calculate_batch_cost_and_usage( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: Optional[str] = None, + model_info: Optional[dict] = None, ) -> Tuple[float, Usage, List[str]]: """ - Calculate the cost and usage of a batch + Calculate the cost and usage of a batch. + + Args: + model_info: Optional deployment-level model info with custom batch + pricing. Threaded through to batch_cost_calculator so that + deployment-specific pricing (e.g. input_cost_per_token_batches) + is used instead of the global cost map. """ batch_cost = _batch_cost_calculator( custom_llm_provider=custom_llm_provider, file_content_dictionary=file_content_dictionary, model_name=model_name, + model_info=model_info, ) batch_usage = _get_batch_job_total_usage_from_file_content( file_content_dictionary=file_content_dictionary, @@ -94,6 +102,7 @@ def _batch_cost_calculator( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", model_name: Optional[str] = None, + model_info: Optional[dict] = None, ) -> float: """ Calculate the cost of a batch based on the output file id @@ -108,6 +117,7 @@ def _batch_cost_calculator( total_cost = _get_batch_job_cost_from_file_content( file_content_dictionary=file_content_dictionary, custom_llm_provider=custom_llm_provider, + model_info=model_info, ) verbose_logger.debug("total_cost=%s", total_cost) return total_cost @@ -290,10 +300,13 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + model_info: Optional[dict] = None, ) -> float: """ Get the cost of a batch job from the file content """ + from litellm.cost_calculator import batch_cost_calculator + try: total_cost: float = 0.0 # parse the file content as json @@ -303,11 +316,22 @@ def _get_batch_job_cost_from_file_content( for _item in file_content_dictionary: if _batch_response_was_successful(_item): _response_body = _get_response_from_batch_job_output_file(_item) - total_cost += litellm.completion_cost( - completion_response=_response_body, - custom_llm_provider=custom_llm_provider, - call_type=CallTypes.aretrieve_batch.value, - ) + if model_info is not None: + usage = _get_batch_job_usage_from_response_body(_response_body) + model = _response_body.get("model", "") + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, + model=model, + custom_llm_provider=custom_llm_provider, + model_info=model_info, + ) + total_cost += prompt_cost + completion_cost + else: + total_cost += litellm.completion_cost( + completion_response=_response_body, + custom_llm_provider=custom_llm_provider, + call_type=CallTypes.aretrieve_batch.value, + ) verbose_logger.debug("total_cost=%s", total_cost) return total_cost except Exception as e: diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 4ea22dbd90f..48dfec2e8c2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1892,9 +1892,16 @@ def batch_cost_calculator( usage: Usage, model: str, custom_llm_provider: Optional[str] = None, + model_info: Optional[dict] = None, ) -> Tuple[float, float]: """ - Calculate the cost of a batch job + Calculate the cost of a batch job. + + Args: + model_info: Optional deployment-level model info containing custom + batch pricing (e.g. input_cost_per_token_batches). When provided, + skips the global litellm.get_model_info() lookup so that + deployment-specific pricing is used. """ _, custom_llm_provider, _, _ = litellm.get_llm_provider( @@ -1907,12 +1914,13 @@ def batch_cost_calculator( custom_llm_provider, ) - try: - model_info: Optional[ModelInfo] = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) - except Exception: - model_info = None + if model_info is None: + try: + model_info = litellm.get_model_info( + model=model, custom_llm_provider=custom_llm_provider + ) + except Exception: + model_info = None if not model_info: return 0.0, 0.0 diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py new file mode 100644 index 00000000000..8bc1bd5a307 --- /dev/null +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -0,0 +1,131 @@ +""" +Test that batch cost calculation uses custom deployment-level pricing +when model_info is provided. + +Reproduces the bug where `input_cost_per_token_batches` / +`output_cost_per_token_batches` set on a proxy deployment's model_info +are ignored by the batch cost pipeline because they are never threaded +through to `batch_cost_calculator`. +""" + +import pytest + +from litellm.batches.batch_utils import ( + _batch_cost_calculator, + _get_batch_job_cost_from_file_content, + calculate_batch_cost_and_usage, +) +from litellm.cost_calculator import batch_cost_calculator +from litellm.types.utils import Usage + + +# --- helpers --- + +def _make_batch_output_line(prompt_tokens: int = 10, completion_tokens: int = 5): + """Return a single successful batch output line (OpenAI JSONL format).""" + return { + "id": "batch_req_1", + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-test", + "object": "chat.completion", + "model": "fake-batch-model", + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + } + ], + }, + }, + "error": None, + } + + +CUSTOM_MODEL_INFO = { + "input_cost_per_token_batches": 0.00125, + "output_cost_per_token_batches": 0.005, +} + + +# --- tests --- + + +def test_batch_cost_calculator_uses_custom_model_info(): + """batch_cost_calculator should use model_info override when provided.""" + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, + model="fake-batch-model", + custom_llm_provider="openai", + model_info=CUSTOM_MODEL_INFO, + ) + + expected_prompt = 10 * 0.00125 + expected_completion = 5 * 0.005 + assert prompt_cost == pytest.approx(expected_prompt), ( + f"Expected prompt cost {expected_prompt}, got {prompt_cost}" + ) + assert completion_cost == pytest.approx(expected_completion), ( + f"Expected completion cost {expected_completion}, got {completion_cost}" + ) + + +def test_get_batch_job_cost_from_file_content_uses_custom_model_info(): + """_get_batch_job_cost_from_file_content should thread model_info to completion_cost.""" + file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] + + cost = _get_batch_job_cost_from_file_content( + file_content_dictionary=file_content, + custom_llm_provider="openai", + model_info=CUSTOM_MODEL_INFO, + ) + + expected = (10 * 0.00125) + (5 * 0.005) + assert cost == pytest.approx(expected), ( + f"Expected total cost {expected}, got {cost}" + ) + + +def test_batch_cost_calculator_func_uses_custom_model_info(): + """_batch_cost_calculator should thread model_info.""" + file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] + + cost = _batch_cost_calculator( + file_content_dictionary=file_content, + custom_llm_provider="openai", + model_info=CUSTOM_MODEL_INFO, + ) + + expected = (10 * 0.00125) + (5 * 0.005) + assert cost == pytest.approx(expected), ( + f"Expected total cost {expected}, got {cost}" + ) + + +@pytest.mark.asyncio +async def test_calculate_batch_cost_and_usage_uses_custom_model_info(): + """calculate_batch_cost_and_usage should thread model_info.""" + file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] + + batch_cost, batch_usage, batch_models = await calculate_batch_cost_and_usage( + file_content_dictionary=file_content, + custom_llm_provider="openai", + model_info=CUSTOM_MODEL_INFO, + ) + + expected = (10 * 0.00125) + (5 * 0.005) + assert batch_cost == pytest.approx(expected), ( + f"Expected total cost {expected}, got {batch_cost}" + ) + assert batch_usage.prompt_tokens == 10 + assert batch_usage.completion_tokens == 5 From 7c3020c04ac30a53b74cc15def86b035052cc20b Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sat, 14 Feb 2026 16:41:35 -0300 Subject: [PATCH 27/50] fix(test): update test_other_constraints_preserved for new schema filtering behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #20813 changed the Anthropic schema filter to remove string and numeric constraints (minLength/maxLength, minimum/maximum) per Anthropic API requirements, but forgot to update the corresponding test. The new behavior (per Anthropic SDK): 1. Remove unsupported constraints from schema (Anthropic API doesn't support them) 2. Add constraint info to description field (e.g., "Note: minimum length: 1") **Changes:** - Updated test to expect constraints REMOVED from schema - Added assertions to verify constraints are added to description - Updated docstring to explain the new behavior **Testing:** - ✅ test_other_constraints_preserved now passes - ✅ All 4 tests in test_anthropic_structured_output.py pass **Related:** - Fixes test broken by PR #20813 - Aligns with Anthropic API requirements documented in commit 84934a7258 Co-Authored-By: Claude Sonnet 4.5 --- .../test_anthropic_structured_output.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py b/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py index 521da6e8c62..705edeaf69c 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_structured_output.py @@ -5,10 +5,9 @@ This test file verifies that Pydantic models with various constraints are properly converted to Anthropic-compatible JSON schemas. """ -from typing import List - import pytest from pydantic import BaseModel, Field +from typing import List class TestAnthropicStructuredOutput: @@ -24,7 +23,7 @@ class TestAnthropicStructuredOutput: Related issue: https://github.com/BerriAI/litellm/issues/19444 """ from litellm.llms.anthropic.chat.transformation import AnthropicConfig - + # Define a Pydantic model with max_length on a List field class ResponseModel(BaseModel): items: List[str] = Field(max_length=5, description="List of items") @@ -131,12 +130,10 @@ class TestAnthropicStructuredOutput: def test_other_constraints_preserved(self): """ - Test that string and numeric constraints are moved to description. + Test that constraints are properly handled (removed from schema, added to description). - Anthropic's output_format API doesn't support minLength/maxLength for - strings or minimum/maximum for numbers. Per Anthropic's SDK approach, - these constraints are removed from the schema and added to the - description text instead. + Per Anthropic API requirements, constraints like minLength/maxLength and + minimum/maximum must be removed from the schema but documented in descriptions. """ from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -150,7 +147,7 @@ class TestAnthropicStructuredOutput: response_format = { "type": "json_schema", - "json_schema": json_schema["json_schema"], + "json_schema": json_schema["json_schema"] } output_format = config.map_response_format_to_anthropic_output_format( @@ -160,16 +157,20 @@ class TestAnthropicStructuredOutput: assert output_format is not None transformed_schema = output_format["schema"] - # String constraints moved to description (not preserved in schema) + # String constraints should be REMOVED from schema (Anthropic doesn't support them) name_schema = transformed_schema["properties"]["name"] assert "maxLength" not in name_schema assert "minLength" not in name_schema - assert "maximum length: 100" in name_schema["description"] + # But constraint info should be added to description + assert "description" in name_schema assert "minimum length: 1" in name_schema["description"] + assert "maximum length: 100" in name_schema["description"] - # Number constraints moved to description (not preserved in schema) + # Number constraints should be REMOVED from schema (Anthropic doesn't support them) age_schema = transformed_schema["properties"]["age"] assert "minimum" not in age_schema assert "maximum" not in age_schema + # But constraint info should be added to description + assert "description" in age_schema assert "minimum value: 0" in age_schema["description"] assert "maximum value: 150" in age_schema["description"] From 03d67d7801c63cea41c46095ab4016b1d7c2396f Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sun, 15 Feb 2026 19:44:23 -0300 Subject: [PATCH 28/50] fix(test): mock vertexai module in GPT-OSS tests to prevent authentication The test_vertex_ai_gpt_oss_simple_request and test_vertex_ai_gpt_oss_reasoning_effort tests were failing in CI with 401 authentication errors. This was because the vertexai module import was triggering authentication attempts even though the _ensure_access_token method was mocked. Added patch.dict('sys.modules', ...) to mock the vertexai module entirely, preventing it from trying to authenticate when imported. This ensures tests are fully isolated and don't attempt real API calls regardless of environment variables or test execution order. This follows the same pattern used in other Vertex AI tests and works in combination with the autouse fixture that clears environment variables. Co-Authored-By: Claude Sonnet 4.5 --- .../test_vertex_ai_gpt_oss_transformation.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py index df1fe26df8b..ed1944b562e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py @@ -73,7 +73,7 @@ class TestVertexAIGPTOSSTransformation: @pytest.mark.asyncio async def test_vertex_ai_gpt_oss_simple_request(): """ - Test that a simple request to vertex_ai/openai/gpt-oss-20b-maas lands at the correct URL + Test that a simple request to vertex_ai/openai/gpt-oss-20b-maas lands at the correct URL with the correct request body. """ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -106,14 +106,20 @@ async def test_vertex_ai_gpt_oss_simple_request(): "total_tokens": 70 } } - + client = AsyncHTTPHandler() - + async def mock_post_func(*args, **kwargs): return mock_response - + + # Mock vertexai module to prevent import from triggering authentication + mock_vertexai = MagicMock() + mock_vertexai.preview = MagicMock() + mock_vertexai.preview.language_models = MagicMock() + with patch.object(client, "post", side_effect=mock_post_func) as mock_post, \ - patch.object(VertexLLM, "_ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718")): + patch.object(VertexLLM, "_ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718")), \ + patch.dict('sys.modules', {'vertexai': mock_vertexai, 'vertexai.preview': mock_vertexai.preview}): response = await litellm.acompletion( model="vertex_ai/openai/gpt-oss-20b-maas", messages=[ @@ -122,7 +128,7 @@ async def test_vertex_ai_gpt_oss_simple_request(): "content": "Your name is Litellm Bot, you are a helpful assistant" }, { - "role": "user", + "role": "user", "content": "Hello, what is your name and can you tell me the weather?" } ], @@ -170,7 +176,7 @@ async def test_vertex_ai_gpt_oss_simple_request(): @pytest.mark.asyncio async def test_vertex_ai_gpt_oss_reasoning_effort(): """ - Test that reasoning_effort parameter is correctly passed in the request body + Test that reasoning_effort parameter is correctly passed in the request body for GPT-OSS models. """ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -184,7 +190,7 @@ async def test_vertex_ai_gpt_oss_reasoning_effort(): mock_response.headers = {} mock_response.json.return_value = { "id": "chatcmpl-test456", - "object": "chat.completion", + "object": "chat.completion", "created": 1234567890, "model": "openai/gpt-oss-20b-maas", "choices": [ @@ -203,14 +209,20 @@ async def test_vertex_ai_gpt_oss_reasoning_effort(): "total_tokens": 67 } } - + client = AsyncHTTPHandler() - + async def mock_post_func(*args, **kwargs): return mock_response - + + # Mock vertexai module to prevent import from triggering authentication + mock_vertexai = MagicMock() + mock_vertexai.preview = MagicMock() + mock_vertexai.preview.language_models = MagicMock() + with patch.object(client, "post", side_effect=mock_post_func) as mock_post, \ - patch.object(VertexLLM, "_ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718")): + patch.object(VertexLLM, "_ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718")), \ + patch.dict('sys.modules', {'vertexai': mock_vertexai, 'vertexai.preview': mock_vertexai.preview}): response = await litellm.acompletion( model="vertex_ai/openai/gpt-oss-20b-maas", messages=[ From 5f79bf4906cbcd2ffbfdeccd6d092e373fb19f0c Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sun, 15 Feb 2026 20:23:27 -0300 Subject: [PATCH 29/50] fix(test): clear tokenizer LRU cache for test isolation The _select_tokenizer_helper function is decorated with @lru_cache, which causes test failures when tests run sequentially with --dist=loadscope. Previous tests' cached results prevent from_pretrained from being called, causing mock assertions to fail. Implemented triple-layer cache clearing: 1. Module-level clear on import 2. Class-level clear in setUpClass 3. Function-level clear in setUp + pytest fixture This ensures test isolation while allowing --dist=loadscope to provide better overall CI stability (70% pass rate vs 40% without loadscope). Fixes the intermittent failure in TestTokenizerSelection where 'from_pretrained' mock was never called due to cache hits. Co-Authored-By: Claude Sonnet 4.5 --- .../litellm_core_utils/test_token_counter.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 875f3db15d5..1506e7b5af0 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -491,7 +491,32 @@ from unittest.mock import MagicMock, patch from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding +# Clear the cache at module load to ensure clean state +_select_tokenizer_helper.cache_clear() + + +@pytest.fixture(autouse=True, scope="function") +def clear_tokenizer_cache_before_test(): + """Clear the LRU cache before each test to ensure test isolation. + + The _select_tokenizer_helper function is decorated with @lru_cache, + which can cause cache hits from previous tests when running with + --dist=loadscope (tests from same file run on same worker). + """ + # Clear before test + _select_tokenizer_helper.cache_clear() + yield + + class TestTokenizerSelection(unittest.TestCase): + @classmethod + def setUpClass(cls): + """Clear cache before class starts.""" + _select_tokenizer_helper.cache_clear() + + def setUp(self): + """Clear cache before each test method.""" + _select_tokenizer_helper.cache_clear() @patch("litellm.utils.Tokenizer.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): # Setup mock to raise an error From bae8816c35c1dfbbcc64f5a732ed9fc16c969de3 Mon Sep 17 00:00:00 2001 From: jquinter Date: Sun, 15 Feb 2026 20:27:29 -0300 Subject: [PATCH 30/50] Update tests/test_litellm/litellm_core_utils/test_token_counter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/litellm_core_utils/test_token_counter.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 1506e7b5af0..85483ba251f 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -517,6 +517,7 @@ class TestTokenizerSelection(unittest.TestCase): def setUp(self): """Clear cache before each test method.""" _select_tokenizer_helper.cache_clear() + @patch("litellm.utils.Tokenizer.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): # Setup mock to raise an error From e82fc28f42249c459320e428bdc0d5082bd14939 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sun, 15 Feb 2026 20:33:09 -0300 Subject: [PATCH 31/50] fix(deps): add fakeredis for pod lock manager tests The test file test_e2e_pod_lock_manager.py requires fakeredis but it was not declared as a dev dependency, causing import errors when the test module is loaded. This is a pre-existing issue that was exposed by better test coverage in PR 21277 but is not caused by that PR. Co-Authored-By: Claude Sonnet 4.5 --- poetry.lock | 58 ++++++++++++++++++++++++++++++++++++++++++-------- pyproject.toml | 1 + 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/poetry.lock b/poetry.lock index 875fcbb3d96..4d44b36aa26 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand. [[package]] name = "a2a-sdk" @@ -330,14 +330,14 @@ zookeeper = ["kazoo"] name = "async-timeout" version = "5.0.1" description = "Timeout context manager for asyncio programs" -optional = true +optional = false python-versions = ">=3.8" -groups = ["main"] -markers = "python_full_version < \"3.11.3\" and (extra == \"extra-proxy\" or extra == \"proxy\") or python_version < \"3.11\"" +groups = ["main", "dev"] files = [ {file = "async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c"}, {file = "async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3"}, ] +markers = {main = "python_full_version < \"3.11.3\" and (extra == \"extra-proxy\" or extra == \"proxy\") or python_version < \"3.11\"", dev = "python_full_version < \"3.11.3\""} [[package]] name = "attrs" @@ -1268,6 +1268,34 @@ typing-extensions = {version = ">=4.6.0", markers = "python_version < \"3.13\""} [package.extras] test = ["pytest (>=6)"] +[[package]] +name = "fakeredis" +version = "2.33.0" +description = "Python implementation of redis API, can be used for testing purposes." +optional = false +python-versions = ">=3.7" +groups = ["dev"] +files = [ + {file = "fakeredis-2.33.0-py3-none-any.whl", hash = "sha256:de535f3f9ccde1c56672ab2fdd6a8efbc4f2619fc2f1acc87b8737177d71c965"}, + {file = "fakeredis-2.33.0.tar.gz", hash = "sha256:d7bc9a69d21df108a6451bbffee23b3eba432c21a654afc7ff2d295428ec5770"}, +] + +[package.dependencies] +redis = [ + {version = ">=4.3", markers = "python_version > \"3.8\""}, + {version = ">=4.3,<7.1.0", markers = "python_version < \"3.10\" and python_version > \"3.8\""}, +] +sortedcontainers = ">=2" +typing-extensions = {version = ">=4.7,<5.0", markers = "python_version < \"3.11\""} + +[package.extras] +bf = ["pyprobables (>=0.6)"] +cf = ["pyprobables (>=0.6)"] +json = ["jsonpath-ng (>=1.6)"] +lua = ["lupa (>=2.1)"] +probabilistic = ["pyprobables (>=0.6)"] +valkey = ["valkey (>=6) ; python_version >= \"3.8\""] + [[package]] name = "fastapi" version = "0.121.3" @@ -5155,7 +5183,7 @@ version = "2.10.1" description = "JSON Web Token implementation in Python" optional = false python-versions = ">=3.9" -groups = ["main", "proxy-dev"] +groups = ["main", "dev", "proxy-dev"] files = [ {file = "PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb"}, {file = "pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953"}, @@ -5534,14 +5562,14 @@ files = [ name = "redis" version = "5.3.1" description = "Python client for Redis database and key-value store" -optional = true +optional = false python-versions = ">=3.8" -groups = ["main"] -markers = "(python_version < \"3.14\" or extra == \"proxy\") and (python_version <= \"3.13\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"extra-proxy\" or extra == \"proxy\")" +groups = ["main", "dev"] files = [ {file = "redis-5.3.1-py3-none-any.whl", hash = "sha256:dc1909bd24669cc31b5f67a039700b16ec30571096c5f1f0d9d2324bff31af97"}, {file = "redis-5.3.1.tar.gz", hash = "sha256:ca49577a531ea64039b5a36db3d6cd1a0c7a60c34124d46924a45b956e8cf14c"}, ] +markers = {main = "(python_version < \"3.14\" or extra == \"proxy\") and (python_version <= \"3.13\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"extra-proxy\" or extra == \"proxy\")"} [package.dependencies] async-timeout = {version = ">=4.0.3", markers = "python_full_version < \"3.11.3\""} @@ -6391,6 +6419,18 @@ files = [ {file = "snowballstemmer-3.0.1.tar.gz", hash = "sha256:6d5eeeec8e9f84d4d56b847692bacf79bc2c8e90c7f80ca4444ff8b6f2e52895"}, ] +[[package]] +name = "sortedcontainers" +version = "2.4.0" +description = "Sorted Containers -- Sorted List, Sorted Dict, Sorted Set" +optional = false +python-versions = "*" +groups = ["dev"] +files = [ + {file = "sortedcontainers-2.4.0-py2.py3-none-any.whl", hash = "sha256:a163dcaede0f1c021485e957a39245190e74249897e2ae4b2aa38595db237ee0"}, + {file = "sortedcontainers-2.4.0.tar.gz", hash = "sha256:25caa5a06cc30b6b83d11423433f65d1f9d76c4c6a0c90e3379eaa43b9bfdb88"}, +] + [[package]] name = "soundfile" version = "0.12.1" @@ -7701,4 +7741,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "0415dbaddd69fc406df167f28db97140b003975dbb2a34d8b594d47fa9feef65" +content-hash = "f2b4f98542c48ba2316a4c90563fc3551f34d9e3771bac39044f55a390e1f1c1" diff --git a/pyproject.toml b/pyproject.toml index d7591560bcc..5726d33c14c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -164,6 +164,7 @@ opentelemetry-sdk = "^1.28.0" opentelemetry-exporter-otlp = "^1.28.0" langfuse = "^2.45.0" fastapi-offline = "^1.7.3" +fakeredis = "^2.27.1" [tool.poetry.group.proxy-dev.dependencies] prisma = "0.11.0" From 706792ba96ba5bd3f48b4e2e80f05556ee93d5ba Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sun, 15 Feb 2026 20:36:24 -0300 Subject: [PATCH 32/50] refactor: simplify cache clearing to avoid over-engineering Based on Greptile feedback: - Removed autouse fixture (applied too broadly to unrelated tests) - Removed setUpClass (redundant since setUp runs before every test) - Kept module-level clear and setUp() method (sufficient for test isolation) - Added blank line for proper formatting The simplified approach still ensures test isolation under --dist=loadscope while avoiding unnecessary complexity. Co-Authored-By: Claude Sonnet 4.5 --- .../litellm_core_utils/test_token_counter.py | 25 +++++-------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 85483ba251f..20f9a6e4279 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -495,27 +495,14 @@ from litellm.utils import _select_tokenizer_helper, claude_json_str, encoding _select_tokenizer_helper.cache_clear() -@pytest.fixture(autouse=True, scope="function") -def clear_tokenizer_cache_before_test(): - """Clear the LRU cache before each test to ensure test isolation. - - The _select_tokenizer_helper function is decorated with @lru_cache, - which can cause cache hits from previous tests when running with - --dist=loadscope (tests from same file run on same worker). - """ - # Clear before test - _select_tokenizer_helper.cache_clear() - yield - - class TestTokenizerSelection(unittest.TestCase): - @classmethod - def setUpClass(cls): - """Clear cache before class starts.""" - _select_tokenizer_helper.cache_clear() - def setUp(self): - """Clear cache before each test method.""" + """Clear the LRU cache before each test method. + + The _select_tokenizer_helper function is decorated with @lru_cache, + which can cause cache hits from previous tests when running with + --dist=loadscope (tests from same file run on same worker). + """ _select_tokenizer_helper.cache_clear() @patch("litellm.utils.Tokenizer.from_pretrained") From 2d41b03f8b66cdb74b58e53d385a2b6b3a42967d Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sun, 15 Feb 2026 20:44:17 -0300 Subject: [PATCH 33/50] fix(test): mock environment variables for callback validation test The test test_proxy_config_state_post_init_callback_call was failing with: ``` ValidationError: 2 validation errors for TeamCallbackMetadata callback_vars.langfuse_public_key Input should be a valid string [type=string_type, input_value=None, input_type=NoneType] ``` Root cause: The test uses environment variable references like "os.environ/LANGFUSE_PUBLIC_KEY" which get resolved at runtime. In parallel execution with --dist=loadscope, these environment variables may not be set in all worker processes, causing the resolution to return None, which fails Pydantic validation expecting strings. Solution: Use monkeypatch to set the required environment variables before the test runs. This ensures consistent behavior across all test execution environments (local, CI, parallel workers). Fixes test failure exposed by PR #21277. Co-Authored-By: Claude Sonnet 4.5 --- tests/proxy_unit_tests/test_proxy_utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 9aecfe9886b..e34cad66ba6 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1170,17 +1170,25 @@ def test_team_callback_metadata_none_values(none_key): assert none_key not in resp -def test_proxy_config_state_post_init_callback_call(): +def test_proxy_config_state_post_init_callback_call(monkeypatch): """ Ensures team_id is still in config, after callback is called Addresses issue: https://github.com/BerriAI/litellm/issues/6787 Where team_id was being popped from config, after callback was called + + Note: Environment variables are mocked to avoid validation errors + in parallel execution where env vars may not be set. """ from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.proxy_server import ProxyConfig + # Mock environment variables to avoid Pydantic validation errors + # when env vars are resolved to None in parallel execution + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "test_public_key") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "test_secret_key") + pc = ProxyConfig() pc.update_config_state( From 01cdec57716a89317eb407d2bea58c7a1d19a8be Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 14:03:31 +0530 Subject: [PATCH 34/50] Fix converse anthropic usage object according to v1/messages specs --- .../adapters/streaming_iterator.py | 9 +- .../adapters/transformation.py | 14 ++- ...al_pass_through_adapters_transformation.py | 105 ++++++++++++++++++ 3 files changed, 125 insertions(+), 3 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index a86820f82e8..de634ff9ecf 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -239,8 +239,13 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): merged_chunk["delta"] = {} # Add usage to the held chunk + uncached_input_tokens = chunk.usage.prompt_tokens or 0 + if hasattr(chunk.usage, "prompt_tokens_details") and chunk.usage.prompt_tokens_details: + cached_tokens = getattr(chunk.usage.prompt_tokens_details, "cached_tokens", 0) or 0 + uncached_input_tokens -= cached_tokens + usage_dict: UsageDelta = { - "input_tokens": chunk.usage.prompt_tokens or 0, + "input_tokens": uncached_input_tokens, "output_tokens": chunk.usage.completion_tokens or 0, } # Add cache tokens if available (for prompt caching support) @@ -412,6 +417,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if block_type == "tool_use": # Type narrowing: content_block_start is ToolUseBlock when block_type is "tool_use" from typing import cast + from litellm.types.llms.anthropic import ToolUseBlock tool_block = cast(ToolUseBlock, content_block_start) @@ -430,6 +436,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # if we get a function name since it signals a new tool call if block_type == "tool_use": from typing import cast + from litellm.types.llms.anthropic import ToolUseBlock tool_block = cast(ToolUseBlock, content_block_start) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 169b138a5f7..efbac13735c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1070,8 +1070,13 @@ class LiteLLMAnthropicMessagesAdapter: ) # extract usage usage: Usage = getattr(response, "usage") + uncached_input_tokens = usage.prompt_tokens or 0 + if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details: + cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0 + uncached_input_tokens -= cached_tokens + anthropic_usage = AnthropicUsage( - input_tokens=usage.prompt_tokens or 0, + input_tokens=uncached_input_tokens, output_tokens=usage.completion_tokens or 0, ) # Add cache tokens if available (for prompt caching support) @@ -1230,8 +1235,13 @@ class LiteLLMAnthropicMessagesAdapter: else: litellm_usage_chunk = None if litellm_usage_chunk is not None: + uncached_input_tokens = litellm_usage_chunk.prompt_tokens or 0 + if hasattr(litellm_usage_chunk, "prompt_tokens_details") and litellm_usage_chunk.prompt_tokens_details: + cached_tokens = getattr(litellm_usage_chunk.prompt_tokens_details, "cached_tokens", 0) or 0 + uncached_input_tokens -= cached_tokens + usage_delta = UsageDelta( - input_tokens=litellm_usage_chunk.prompt_tokens or 0, + input_tokens=uncached_input_tokens, output_tokens=litellm_usage_chunk.completion_tokens or 0, ) # Add cache tokens if available (for prompt caching support) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index b228a51447b..f9e5c6d0252 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1706,3 +1706,108 @@ def test_translate_openai_response_restores_tool_names(): assert len(tool_use_blocks) == 1 # Name should be restored to original assert tool_use_blocks[0]["name"] == original_name + + +def test_translate_openai_response_to_anthropic_input_tokens_excludes_cached_tokens(): + """ + Regression test: input_tokens in Anthropic format should NOT include cached tokens. + + Issue: v1/messages API was returning incorrect input_token count when using prompt caching. + The OpenAI format includes cached tokens in prompt_tokens, but Anthropic format should not. + + According to Anthropic's spec: + - input_tokens = uncached input tokens only + - cache_read_input_tokens = tokens read from cache + + In OpenAI format: + - prompt_tokens = all input tokens (including cached) + - prompt_tokens_details.cached_tokens = cached tokens + + Expected: anthropic.input_tokens = openai.prompt_tokens - openai.prompt_tokens_details.cached_tokens + """ + from litellm.types.utils import PromptTokensDetailsWrapper + + # Create OpenAI format response with cached tokens + # Scenario: 100 total prompt tokens, 30 of which are cached + usage = Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=30 + ), + cache_read_input_tokens=30, # Anthropic format cache info + ) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + + # Convert to Anthropic format + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + # Validate: input_tokens should be 70 (100 - 30 cached), not 100 + assert anthropic_response["usage"]["input_tokens"] == 70, ( + f"Expected input_tokens=70 (100 total - 30 cached), " + f"but got {anthropic_response['usage']['input_tokens']}. " + f"input_tokens should NOT include cached tokens per Anthropic spec." + ) + assert anthropic_response["usage"]["output_tokens"] == 50 + assert anthropic_response["usage"]["cache_read_input_tokens"] == 30 + + +def test_translate_openai_response_to_anthropic_input_tokens_no_cache(): + """ + Regression test: input_tokens should equal prompt_tokens when there are no cached tokens. + """ + from litellm.types.utils import PromptTokensDetailsWrapper + + # Create OpenAI format response without cached tokens + usage = Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + ) + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message( + role="assistant", + content="Test response", + ), + ) + ], + model="claude-3-sonnet-20240229", + usage=usage, + ) + + # Convert to Anthropic format + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_response = adapter.translate_openai_response_to_anthropic( + response=response, + tool_name_mapping=None, + ) + + # Validate: input_tokens should equal prompt_tokens when no caching + assert anthropic_response["usage"]["input_tokens"] == 100 + assert anthropic_response["usage"]["output_tokens"] == 50 From ee9e997755c8fdb01aa171be101ebed240496598 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 17:02:34 +0530 Subject: [PATCH 35/50] Add routing based on if reasoning is supported or not --- .../experimental_pass_through/adapters/handler.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index c6caaddf98b..73e74c228ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -19,6 +19,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.utils import ModelResponse +from litellm.utils import get_model_info if TYPE_CHECKING: pass @@ -63,6 +64,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: return model = completion_kwargs.get("model") + try: + model_info = get_model_info(model=cast(str, model), custom_llm_provider=custom_llm_provider) + if model_info and model_info.get("supports_reasoning") is False: + # Model doesn't support reasoning/responses API, don't route + return + except Exception: + pass + if isinstance(model, str) and model and not model.startswith("responses/"): # Prefix model with "responses/" to route to OpenAI Responses API completion_kwargs["model"] = f"responses/{model}" From c00c5a6e29dc7d54290d38aa36183abd3f5664c3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 17:39:04 +0530 Subject: [PATCH 36/50] add fireworks_ai/accounts/fireworks/models/kimi-k2p5 in model map --- litellm/model_prices_and_context_window_backup.json | 13 +++++++++++++ model_prices_and_context_window.json | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 95d8ba2ff60..13f49e04985 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12456,6 +12456,19 @@ "supports_tool_choice": true, "supports_web_search": true }, + "fireworks_ai/accounts/fireworks/models/kimi-k2p5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://fireworks.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "fireworks_ai/accounts/fireworks/models/llama-v3p1-405b-instruct": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 95d8ba2ff60..13f49e04985 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12456,6 +12456,19 @@ "supports_tool_choice": true, "supports_web_search": true }, + "fireworks_ai/accounts/fireworks/models/kimi-k2p5": { + "input_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://fireworks.ai/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "fireworks_ai/accounts/fireworks/models/llama-v3p1-405b-instruct": { "input_cost_per_token": 3e-06, "litellm_provider": "fireworks_ai", From cc4e022176d6546256397b98b22dc147e0d9d1b8 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Mon, 16 Feb 2026 07:14:42 -0500 Subject: [PATCH 37/50] Removed stray .md file --- fix_afile_retrieve_returns_unified_id.md | 47 ------------------------ 1 file changed, 47 deletions(-) delete mode 100644 fix_afile_retrieve_returns_unified_id.md diff --git a/fix_afile_retrieve_returns_unified_id.md b/fix_afile_retrieve_returns_unified_id.md deleted file mode 100644 index c8cc72d6f99..00000000000 --- a/fix_afile_retrieve_returns_unified_id.md +++ /dev/null @@ -1,47 +0,0 @@ -# Fix: afile_retrieve returns raw provider ID for batch output files - -## Bug - -`managed_files.afile_retrieve()` Case 2 (file_object already in DB) returned the stored `file_object` without replacing `.id` with the unified file ID. Case 3 (fetch from provider) did this correctly at line 1028. - -## Fix - -One-line change in `enterprise/litellm_enterprise/proxy/hooks/managed_files.py`: - -```python -# Before (line 1013-1014) -if stored_file_object and stored_file_object.file_object: - return stored_file_object.file_object - -# After -if stored_file_object and stored_file_object.file_object: - stored_file_object.file_object.id = file_id - return stored_file_object.file_object -``` - -## Test - -```bash -poetry run pytest tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py -s -vvvv -``` - -## Test failure (before fix) - -``` -FAILED tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py::test_should_return_unified_id_when_file_object_exists_in_db -AssertionError: afile_retrieve should return the unified ID 'bGl0ZWxsbV9wcm94eTp1bmlmaWVkX291dHB1dF9maWxl', but got raw provider ID 'batch_20260214-output-file-1' -assert 'batch_20260214-output-file-1' == 'bGl0ZWxsbV9wcm94eTp1bmlmaWVkX291dHB1dF9maWxl' -=================== 1 failed, 1 retried in 102.95s =================== -``` - -## Test pass (after fix) - -``` -tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py::test_should_return_unified_id_when_file_object_exists_in_db PASSED -============================== 1 passed in 0.11s =============================== -``` - -## Files changed - -- `enterprise/litellm_enterprise/proxy/hooks/managed_files.py` — one-line fix -- `tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py` — new test From 37da38fdaa999dea1256329fa608c1b2cd9881c0 Mon Sep 17 00:00:00 2001 From: mjkam Date: Mon, 16 Feb 2026 21:49:41 +0900 Subject: [PATCH 38/50] fix(bedrock): clamp thinking.budget_tokens to minimum 1024 Bedrock rejects thinking.budget_tokens values below 1024 with a 400 error. This adds automatic clamping in the LiteLLM transformation layer so callers (e.g. router with reasoning_effort="low") don't need to know about the provider-specific minimum. Fixes #21297 Co-Authored-By: Claude Opus 4.6 --- litellm/constants.py | 3 ++ .../bedrock/chat/converse_transformation.py | 29 +++++++++++- .../chat/test_converse_transformation.py | 44 +++++++++++++++++++ 3 files changed, 75 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 03f80a8cb78..f83b78ce6bf 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -319,6 +319,9 @@ NON_LLM_CONNECTION_TIMEOUT = int( MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000)) MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048)) BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75)) +BEDROCK_MIN_THINKING_BUDGET_TOKENS = int( + os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024) +) REPLICATE_POLLING_DELAY_SECONDS = float( os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5) ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index efa755d515e..5faae07e2b9 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -11,7 +11,10 @@ import httpx import litellm from litellm._logging import verbose_logger -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import ( + BEDROCK_MIN_THINKING_BUDGET_TOKENS, + RESPONSE_FORMAT_TOOL_NAME, +) from litellm.litellm_core_utils.core_helpers import ( filter_exceptions_from_params, filter_internal_params, @@ -434,6 +437,25 @@ class AmazonConverseConfig(BaseConfig): reasoning_effort=reasoning_effort, model=model ) + @staticmethod + def _clamp_thinking_budget_tokens(optional_params: dict) -> None: + """ + Clamp thinking.budget_tokens to the Bedrock minimum (1024). + + Bedrock returns a 400 error if budget_tokens < 1024. + """ + thinking = optional_params.get("thinking") + if isinstance(thinking, dict): + budget = thinking.get("budget_tokens") + if isinstance(budget, int) and budget < BEDROCK_MIN_THINKING_BUDGET_TOKENS: + verbose_logger.debug( + "Bedrock requires thinking.budget_tokens >= %d, got %d. " + "Clamping to minimum.", + BEDROCK_MIN_THINKING_BUDGET_TOKENS, + budget, + ) + thinking["budget_tokens"] = BEDROCK_MIN_THINKING_BUDGET_TOKENS + def get_supported_openai_params(self, model: str) -> List[str]: from litellm.utils import supports_function_calling @@ -871,9 +893,14 @@ class AmazonConverseConfig(BaseConfig): Checks 'non_default_params' for 'thinking' and 'max_tokens' if 'thinking' is enabled and 'max_tokens' is not specified, set 'max_tokens' to the thinking token budget + DEFAULT_MAX_TOKENS + + Also clamps thinking.budget_tokens to the Bedrock minimum (1024) to + prevent 400 errors from the Bedrock API. """ from litellm.constants import DEFAULT_MAX_TOKENS + self._clamp_thinking_budget_tokens(optional_params) + is_thinking_enabled = self.is_thinking_enabled(optional_params) is_max_tokens_in_request = self.is_max_tokens_in_request(non_default_params) if is_thinking_enabled and not is_max_tokens_in_request: diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index ce43f22d8f8..ddbb0454cac 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2934,3 +2934,47 @@ def test_drop_thinking_param_when_thinking_blocks_missing(): finally: # Restore original modify_params setting litellm.modify_params = original_modify_params + + +class TestBedrockMinThinkingBudgetTokens: + """Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024).""" + + def _map_params( + self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0" + ): + """Helper to call map_openai_params with the given thinking value.""" + config = AmazonConverseConfig() + non_default_params = {"thinking": thinking_value} + optional_params = {"thinking": thinking_value} + return config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + def test_budget_tokens_below_minimum_is_clamped(self): + """budget_tokens < 1024 should be clamped to 1024.""" + result = self._map_params({"type": "enabled", "budget_tokens": 499}) + assert result["thinking"]["budget_tokens"] == 1024 + + def test_budget_tokens_at_minimum_is_unchanged(self): + """budget_tokens == 1024 should remain 1024.""" + result = self._map_params({"type": "enabled", "budget_tokens": 1024}) + assert result["thinking"]["budget_tokens"] == 1024 + + def test_budget_tokens_above_minimum_is_unchanged(self): + """budget_tokens > 1024 should remain unchanged.""" + result = self._map_params({"type": "enabled", "budget_tokens": 2048}) + assert result["thinking"]["budget_tokens"] == 2048 + + def test_no_thinking_param_does_not_error(self): + """When thinking is not provided, map_openai_params should not raise.""" + config = AmazonConverseConfig() + result = config.map_openai_params( + non_default_params={}, + optional_params={}, + model="anthropic.claude-3-7-sonnet-20250219-v1:0", + drop_params=False, + ) + assert "thinking" not in result or result.get("thinking") is None From 43f9a588d9b2ba094b12efcb8aa5f59582f85956 Mon Sep 17 00:00:00 2001 From: jquinter Date: Sat, 14 Feb 2026 03:54:50 -0300 Subject: [PATCH 39/50] fix: improve Langfuse test isolation to prevent flaky failures (#21093) The test was creating fresh mocks but not fully isolating from setUp state, causing intermittent CI failures with 'Expected generation to be called once. Called 0 times.' Instead of creating fresh mocks, properly reset the existing setUp mocks to ensure clean state while maintaining proper mock chain configuration. --- .../integrations/test_langfuse.py | 37 +++++++++---------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index cd3d5b9ebe3..010d9f863c2 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -268,22 +268,21 @@ class TestLangfuseUsageDetails(unittest.TestCase): Test that _log_langfuse_v2 correctly handles None values in the usage object by converting them to 0, preventing validation errors. """ - # Create fresh mocks for this test to avoid state pollution from setUp's side_effect - # The setUp configures trace.side_effect which can interfere with return_value - mock_trace = MagicMock() - mock_generation = MagicMock() - mock_generation.trace_id = "test-trace-id" + # Reset the mock to ensure clean state + self.mock_langfuse_client.reset_mock() + self.mock_langfuse_trace.reset_mock() + self.mock_langfuse_generation.reset_mock() + + # Re-setup the trace and generation chain with clean state + self.mock_langfuse_generation.trace_id = "test-trace-id" mock_span = MagicMock() mock_span.end = MagicMock() - - mock_trace.generation.return_value = mock_generation - mock_trace.span.return_value = mock_span - - mock_client = MagicMock() - mock_client.trace.return_value = mock_trace - - # Use our fresh mock client - self.logger.Langfuse = mock_client + self.mock_langfuse_trace.span.return_value = mock_span + self.mock_langfuse_trace.generation.return_value = self.mock_langfuse_generation + + # Ensure trace returns our mock + self.mock_langfuse_client.trace.return_value = self.mock_langfuse_trace + self.logger.Langfuse = self.mock_langfuse_client with patch( "litellm.integrations.langfuse.langfuse._add_prompt_to_generation_params", @@ -337,13 +336,13 @@ class TestLangfuseUsageDetails(unittest.TestCase): ) except Exception as e: self.fail(f"_log_langfuse_v2 raised an exception: {e}") - + # Verify that trace was called first - mock_client.trace.assert_called() - + self.mock_langfuse_client.trace.assert_called() + # Check the arguments passed to the mocked langfuse generation call - mock_trace.generation.assert_called_once() - call_args, call_kwargs = mock_trace.generation.call_args + self.mock_langfuse_trace.generation.assert_called_once() + call_args, call_kwargs = self.mock_langfuse_trace.generation.call_args # Inspect the usage and usage_details dictionaries usage_arg = call_kwargs.get("usage") From 9c71d8b61b3e35062d65e2b315f0f1833138eba9 Mon Sep 17 00:00:00 2001 From: Fly <48186978+tuzkiyoung@users.noreply.github.com> Date: Sat, 14 Feb 2026 14:55:42 +0800 Subject: [PATCH 40/50] feat(s3): add support for virtual-hosted-style URLs (#21094) Add s3_use_virtual_hosted_style parameter to support AWS S3 virtual-hosted-style URL format (bucket.endpoint/key) alongside the existing path-style format (endpoint/bucket/key). This enables compatibility with S3-compatible services like MinIO and aligns with AWS S3 official terminology. --- docs/my-website/docs/proxy/logging.md | 1 + litellm/integrations/s3_v2.py | 75 +++++++--- tests/test_litellm/integrations/test_s3_v2.py | 135 ++++++++++++++++++ 3 files changed, 188 insertions(+), 23 deletions(-) diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 56fb420e6cf..1abb127dfda 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -1338,6 +1338,7 @@ litellm_settings: s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 s3_path: my-test-path # [OPTIONAL] set path in bucket you want to write logs to s3_endpoint_url: https://s3.amazonaws.com # [OPTIONAL] S3 endpoint URL, if you want to use Backblaze/cloudflare s3 buckets + s3_use_virtual_hosted_style: false # [OPTIONAL] use virtual-hosted-style URLs (bucket.endpoint/key) instead of path-style (endpoint/bucket/key). Useful for S3-compatible services like MinIO s3_strip_base64_files: false # [OPTIONAL] remove base64 files before storing in s3 ``` diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 534b85e4752..50278dad5a8 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -51,6 +51,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_team_prefix: bool = False, s3_strip_base64_files: bool = False, s3_use_key_prefix: bool = False, + s3_use_virtual_hosted_style: bool = False, **kwargs, ): try: @@ -78,7 +79,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_path=s3_path, s3_use_team_prefix=s3_use_team_prefix, s3_strip_base64_files=s3_strip_base64_files, - s3_use_key_prefix=s3_use_key_prefix + s3_use_key_prefix=s3_use_key_prefix, + s3_use_virtual_hosted_style=s3_use_virtual_hosted_style ) verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") @@ -135,6 +137,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_team_prefix: bool = False, s3_strip_base64_files: bool = False, s3_use_key_prefix: bool = False, + s3_use_virtual_hosted_style: bool = False, ): """ Initialize the s3 params for this logging callback @@ -217,6 +220,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): or s3_strip_base64_files ) + self.s3_use_virtual_hosted_style = ( + bool(litellm.s3_callback_params.get("s3_use_virtual_hosted_style", False)) + or s3_use_virtual_hosted_style + ) + return async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -302,13 +310,20 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" if self.s3_endpoint_url and self.s3_bucket_name: - url = ( - self.s3_endpoint_url - + "/" - + self.s3_bucket_name - + "/" - + batch_logging_element.s3_object_key - ) + if self.s3_use_virtual_hosted_style: + # Virtual-hosted-style: bucket.endpoint/key + endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" + else: + # Path-style: endpoint/bucket/key + url = ( + self.s3_endpoint_url + + "/" + + self.s3_bucket_name + + "/" + + batch_logging_element.s3_object_key + ) # Convert JSON to string json_string = safe_dumps(batch_logging_element.payload) @@ -456,13 +471,20 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" if self.s3_endpoint_url and self.s3_bucket_name: - url = ( - self.s3_endpoint_url - + "/" - + self.s3_bucket_name - + "/" - + batch_logging_element.s3_object_key - ) + if self.s3_use_virtual_hosted_style: + # Virtual-hosted-style: bucket.endpoint/key + endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" + else: + # Path-style: endpoint/bucket/key + url = ( + self.s3_endpoint_url + + "/" + + self.s3_bucket_name + + "/" + + batch_logging_element.s3_object_key + ) # Convert JSON to string json_string = safe_dumps(batch_logging_element.payload) @@ -550,13 +572,20 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{s3_object_key}" if self.s3_endpoint_url and self.s3_bucket_name: - url = ( - self.s3_endpoint_url - + "/" - + self.s3_bucket_name - + "/" - + s3_object_key - ) + if self.s3_use_virtual_hosted_style: + # Virtual-hosted-style: bucket.endpoint/key + endpoint_host = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{s3_object_key}" + else: + # Path-style: endpoint/bucket/key + url = ( + self.s3_endpoint_url + + "/" + + self.s3_bucket_name + + "/" + + s3_object_key + ) # Prepare the request for GET operation # For GET requests, we need x-amz-content-sha256 with hash of empty string @@ -618,4 +647,4 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.exception( f"Error retrieving object {object_key} from cold storage: {str(e)}" ) - return None + return None \ No newline at end of file diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 0a3523699a9..3a96237f4e3 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -157,6 +157,141 @@ class TestS3V2UnitTests: assert result == {"downloaded": "data"} + @patch('asyncio.create_task') + @patch('litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush') + def test_s3_v2_virtual_hosted_style(self, mock_periodic_flush, mock_create_task): + """Test s3_use_virtual_hosted_style parameter for virtual-hosted-style URLs""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + # Mock periodic_flush and create_task to prevent async task creation during init + mock_periodic_flush.return_value = None + mock_create_task.return_value = None + + # Mock response for all tests + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + + # Create a test batch logging element + test_element = s3BatchLoggingElement( + s3_object_key="2025-09-14/test-key.json", + payload={"test": "data"}, + s3_object_download_filename="test-file.json" + ) + + # Test 1: Virtual-hosted-style with custom endpoint + s3_logger_virtual = S3Logger( + s3_bucket_name="test-bucket", + s3_endpoint_url="https://s3.custom-endpoint.com", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_use_virtual_hosted_style=True + ) + + s3_logger_virtual.async_httpx_client = AsyncMock() + s3_logger_virtual.async_httpx_client.put.return_value = mock_response + + asyncio.run(s3_logger_virtual.async_upload_data_to_s3(test_element)) + + call_args = s3_logger_virtual.async_httpx_client.put.call_args + assert call_args is not None + url = call_args[0][0] + expected_url = "https://test-bucket.s3.custom-endpoint.com/2025-09-14/test-key.json" + assert url == expected_url, f"Expected virtual-hosted-style URL {expected_url}, got {url}" + + # Test 2: Path-style (default behavior with s3_use_virtual_hosted_style=False) + s3_logger_path = S3Logger( + s3_bucket_name="test-bucket", + s3_endpoint_url="https://s3.custom-endpoint.com", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_use_virtual_hosted_style=False + ) + + s3_logger_path.async_httpx_client = AsyncMock() + s3_logger_path.async_httpx_client.put.return_value = mock_response + + asyncio.run(s3_logger_path.async_upload_data_to_s3(test_element)) + + call_args_path = s3_logger_path.async_httpx_client.put.call_args + assert call_args_path is not None + url_path = call_args_path[0][0] + expected_path_url = "https://s3.custom-endpoint.com/test-bucket/2025-09-14/test-key.json" + assert url_path == expected_path_url, f"Expected path-style URL {expected_path_url}, got {url_path}" + + # Test 3: Virtual-hosted-style with http protocol + s3_logger_http = S3Logger( + s3_bucket_name="http-bucket", + s3_endpoint_url="http://minio.local:9000", + s3_aws_access_key_id="minio-key", + s3_aws_secret_access_key="minio-secret", + s3_region_name="us-east-1", + s3_use_virtual_hosted_style=True + ) + + s3_logger_http.async_httpx_client = AsyncMock() + s3_logger_http.async_httpx_client.put.return_value = mock_response + + asyncio.run(s3_logger_http.async_upload_data_to_s3(test_element)) + + call_args_http = s3_logger_http.async_httpx_client.put.call_args + assert call_args_http is not None + url_http = call_args_http[0][0] + expected_http_url = "http://http-bucket.minio.local:9000/2025-09-14/test-key.json" + assert url_http == expected_http_url, f"Expected virtual-hosted-style URL with http {expected_http_url}, got {url_http}" + + # Test 4: Sync upload method with virtual-hosted-style + s3_logger_sync_virtual = S3Logger( + s3_bucket_name="sync-bucket", + s3_endpoint_url="https://storage.example.com", + s3_aws_access_key_id="sync-key", + s3_aws_secret_access_key="sync-secret", + s3_region_name="us-east-1", + s3_use_virtual_hosted_style=True + ) + + mock_sync_client = MagicMock() + mock_sync_client.put.return_value = mock_response + + with patch('litellm.integrations.s3_v2._get_httpx_client', return_value=mock_sync_client): + s3_logger_sync_virtual.upload_data_to_s3(test_element) + + call_args_sync = mock_sync_client.put.call_args + assert call_args_sync is not None + url_sync = call_args_sync[0][0] + expected_sync_url = "https://sync-bucket.storage.example.com/2025-09-14/test-key.json" + assert url_sync == expected_sync_url, f"Expected virtual-hosted-style sync URL {expected_sync_url}, got {url_sync}" + + # Test 5: Download method with virtual-hosted-style + s3_logger_download_virtual = S3Logger( + s3_bucket_name="download-bucket", + s3_endpoint_url="https://download.endpoint.com", + s3_aws_access_key_id="download-key", + s3_aws_secret_access_key="download-secret", + s3_region_name="us-east-1", + s3_use_virtual_hosted_style=True + ) + + mock_download_response = MagicMock() + mock_download_response.status_code = 200 + mock_download_response.json = MagicMock(return_value={"downloaded": "data"}) + s3_logger_download_virtual.async_httpx_client = AsyncMock() + s3_logger_download_virtual.async_httpx_client.get.return_value = mock_download_response + + result = asyncio.run(s3_logger_download_virtual._download_object_from_s3("2025-09-14/download-test-key.json")) + + call_args_download = s3_logger_download_virtual.async_httpx_client.get.call_args + assert call_args_download is not None + url_download = call_args_download[0][0] + expected_download_url = "https://download-bucket.download.endpoint.com/2025-09-14/download-test-key.json" + assert url_download == expected_download_url, f"Expected virtual-hosted-style download URL {expected_download_url}, got {url_download}" + + assert result == {"downloaded": "data"} + @pytest.mark.asyncio async def test_strip_base64_removes_file_and_nontext_entries(): logger = S3Logger(s3_strip_base64_files=True) From a3762e7d4986b5e7d65c0992f39af041489e88e1 Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Mon, 16 Feb 2026 07:58:04 -0500 Subject: [PATCH 41/50] Addressed greptile comments to extract common helpers and return 404 --- .../proxy/hooks/managed_files.py | 13 +++--- litellm/proxy/batches_endpoints/endpoints.py | 41 ++++--------------- .../openai_files_endpoints/common_utils.py | 32 +++++++++++---- .../test_deleted_file_returns_403_not_404.py | 21 ++++------ 4 files changed, 47 insertions(+), 60 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index d9514bf28bc..f341a1e9634 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -230,14 +230,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if managed_file: return managed_file.created_by == user_id - # When DB record is missing (file was deleted), allow through so downstream returns 404. - # Matches can_user_call_unified_object_id which also returns True for missing records. - return True + raise HTTPException( + status_code=404, + detail=f"File not found: {unified_file_id}", + ) async def can_user_call_unified_object_id( self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth ) -> bool: - ## check if the user has access to the unified object id ## check if the user has access to the unified object id user_id = user_api_key_dict.user_id managed_object = ( @@ -248,7 +248,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if managed_object: return managed_object.created_by == user_id - return True # don't raise error if managed object is not found + raise HTTPException( + status_code=404, + detail=f"Object not found: {unified_object_id}", + ) async def list_user_batches( self, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index b125b90338b..783ac9d6f19 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_models_from_unified_file_id, get_original_file_id, prepare_data_with_credentials, + resolve_input_file_id_to_unified, update_batch_in_database, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model @@ -379,22 +380,9 @@ async def retrieve_batch( ) # async_post_call_success_hook replaces batch.id and output_file_id with unified IDs - # but not input_file_id. Look up the unified ID from flat_model_file_ids. - if ( - unified_batch_id - and hasattr(response, "input_file_id") - and response.input_file_id - and not _is_base64_encoded_unified_file_id(response.input_file_id) - and prisma_client - ): - try: - _managed_file = await prisma_client.db.litellm_managedfiletable.find_first( - where={"flat_model_file_ids": {"has": response.input_file_id}} - ) - if _managed_file: - response.input_file_id = _managed_file.unified_file_id - except Exception: - pass + # but not input_file_id. Resolve raw provider ID to unified ID. + if unified_batch_id: + await resolve_input_file_id_to_unified(response, prisma_client) asyncio.create_task( proxy_logging_obj.update_request_status( @@ -497,23 +485,10 @@ async def retrieve_batch( data=data, user_api_key_dict=user_api_key_dict, response=response ) - # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id (terminal state path) - # Same as above — resolve raw provider input_file_id to unified ID. - if ( - unified_batch_id - and hasattr(response, "input_file_id") - and response.input_file_id - and not _is_base64_encoded_unified_file_id(response.input_file_id) - and prisma_client - ): - try: - _managed_file = await prisma_client.db.litellm_managedfiletable.find_first( - where={"flat_model_file_ids": {"has": response.input_file_id}} - ) - if _managed_file: - response.input_file_id = _managed_file.unified_file_id - except Exception: - pass + # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id + # Resolve raw provider input_file_id to unified ID. + if unified_batch_id: + await resolve_input_file_id_to_unified(response, prisma_client) ### ALERTING ### asyncio.create_task( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 0e3b31b2aa7..75f64cddf59 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -644,6 +644,28 @@ def _extract_model_param(request: "Request", request_body: dict) -> Optional[str # ============================================================================ +async def resolve_input_file_id_to_unified(response, prisma_client) -> None: + """ + If the batch response contains a raw provider input_file_id (not already a + unified ID), look up the corresponding unified file ID from the managed file + table and replace it in-place. + """ + if ( + hasattr(response, "input_file_id") + and response.input_file_id + and not _is_base64_encoded_unified_file_id(response.input_file_id) + and prisma_client + ): + try: + managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if managed_file: + response.input_file_id = managed_file.unified_file_id + except Exception: + pass + + async def get_batch_from_database( batch_id: str, unified_batch_id: Union[str, Literal[False]], @@ -689,15 +711,7 @@ async def get_batch_from_database( response.id = batch_id # The stored batch object has the raw provider input_file_id. Resolve to unified ID. - if response.input_file_id and not _is_base64_encoded_unified_file_id(response.input_file_id): - try: - managed_file = await prisma_client.db.litellm_managedfiletable.find_first( - where={"flat_model_file_ids": {"has": response.input_file_id}} - ) - if managed_file: - response.input_file_id = managed_file.unified_file_id - except Exception: - pass + await resolve_input_file_id_to_unified(response, prisma_client) verbose_proxy_logger.debug( f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}" diff --git a/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py b/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py index 042796c9640..7ad564dc8f9 100644 --- a/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py +++ b/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py @@ -1,13 +1,9 @@ """ -Regression test: deleted managed files should not return 403. +Regression test: deleted managed files should return 404, not 403. When a managed file's DB record has been deleted, can_user_call_unified_file_id() -returns False (record not found → treated as "access denied"). This causes -check_managed_file_id_access() to raise 403 instead of allowing the request -through so downstream code can return a proper 404. - -The equivalent method for objects (can_user_call_unified_object_id) already -returns True when the record is missing. +raises HTTPException(404) directly — rather than returning True (which would +weaken access control) or False (which would cause a misleading 403). """ import base64 @@ -49,20 +45,19 @@ def _make_managed_files_with_no_db_record(): @pytest.mark.asyncio -async def test_should_not_raise_403_for_deleted_file(): +async def test_should_raise_404_for_deleted_file(): """ When a managed file record has been deleted from the DB, - check_managed_file_id_access should NOT raise 403. - It should allow the request through so downstream can return 404. + check_managed_file_id_access should raise 404 (not 403). """ unified_file_id = _make_unified_file_id() managed_files = _make_managed_files_with_no_db_record() user = _make_user_api_key_dict("any-user") data = {"file_id": unified_file_id} - # This should NOT raise — deleted file should pass through access check - result = await managed_files.check_managed_file_id_access(data, user) - assert result is True + with pytest.raises(HTTPException) as exc_info: + await managed_files.check_managed_file_id_access(data, user) + assert exc_info.value.status_code == 404 @pytest.mark.asyncio From 51b1b0339c654bf6cc36c132de40235b91624cb1 Mon Sep 17 00:00:00 2001 From: Kristoffer Arlind <13228507+KristofferArlind@users.noreply.github.com> Date: Sat, 14 Feb 2026 15:13:45 +0800 Subject: [PATCH 42/50] Allow effort="max" for Claude Opus 4.6 (#21112) --- litellm/llms/anthropic/chat/transformation.py | 8 +++-- .../test_anthropic_chat_transformation.py | 35 +++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index c2cfff80685..85a4790a9b9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1282,9 +1282,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): output_config = optional_params.get("output_config") if output_config and isinstance(output_config, dict): effort = output_config.get("effort") - if effort and effort not in ["high", "medium", "low"]: + if effort and effort not in ["high", "medium", "low", "max"]: raise ValueError( - f"Invalid effort value: {effort}. Must be one of: 'high', 'medium', 'low'" + f"Invalid effort value: {effort}. Must be one of: 'high', 'medium', 'low', 'max'" + ) + if effort == "max" and not self._is_claude_opus_4_6(model): + raise ValueError( + f"effort='max' is only supported by Claude Opus 4.6. Got model: {model}" ) data["output_config"] = output_config diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 57e0dd494e0..50e948c1a27 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1638,6 +1638,41 @@ def test_effort_with_claude_opus_45(): assert result["model"] == "claude-opus-4-5-20251101" +def test_effort_validation_with_opus_46(): + """Test that all four effort levels are accepted for Claude Opus 4.6.""" + config = AnthropicConfig() + + messages = [{"role": "user", "content": "Test"}] + + for effort in ["high", "medium", "low", "max"]: + optional_params = {"output_config": {"effort": effort}} + result = config.transform_request( + model="claude-opus-4-6-20260205", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + assert result["output_config"]["effort"] == effort + + +def test_max_effort_rejected_for_opus_45(): + """Test that effort='max' is rejected when using Claude Opus 4.5.""" + config = AnthropicConfig() + + messages = [{"role": "user", "content": "Test"}] + + with pytest.raises(ValueError, match="effort='max' is only supported by Claude Opus 4.6"): + optional_params = {"output_config": {"effort": "max"}} + config.transform_request( + model="claude-opus-4-5-20251101", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + + def test_effort_with_other_features(): """Test effort works alongside other features (thinking, tools).""" config = AnthropicConfig() From 7ef908381223a045cc6851b6c10e7ec876d40092 Mon Sep 17 00:00:00 2001 From: Constantine Date: Sat, 14 Feb 2026 10:15:15 +0300 Subject: [PATCH 43/50] fix(aiohttp): prevent closing shared ClientSession in AiohttpTransport (#21117) When a shared ClientSession is passed to LiteLLMAiohttpTransport, calling aclose() on the transport would close the shared session, breaking other clients still using it. Add owns_session parameter (default True for backwards compatibility) to AiohttpTransport and LiteLLMAiohttpTransport. When a shared session is provided in http_handler.py, owns_session=False is set to prevent the transport from closing a session it does not own. This aligns AiohttpTransport with the ownership pattern already used in AiohttpHandler (aiohttp_handler.py). --- .../llms/custom_httpx/aiohttp_transport.py | 12 +++++-- litellm/llms/custom_httpx/http_handler.py | 1 + .../custom_httpx/test_aiohttp_transport.py | 32 +++++++++++++++++++ 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index fb98006c7e4..6cec1f4fe16 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -119,8 +119,13 @@ class AiohttpResponseStream(httpx.AsyncByteStream): class AiohttpTransport(httpx.AsyncBaseTransport): - def __init__(self, client: Union[ClientSession, Callable[[], ClientSession]]) -> None: + def __init__( + self, + client: Union[ClientSession, Callable[[], ClientSession]], + owns_session: bool = True, + ) -> None: self.client = client + self._owns_session = owns_session ######################################################### # Class variables for proxy settings @@ -128,7 +133,7 @@ class AiohttpTransport(httpx.AsyncBaseTransport): self.proxy_cache: Dict[str, Optional[str]] = {} async def aclose(self) -> None: - if isinstance(self.client, ClientSession): + if self._owns_session and isinstance(self.client, ClientSession): await self.client.close() @@ -144,10 +149,11 @@ class LiteLLMAiohttpTransport(AiohttpTransport): self, client: Union[ClientSession, Callable[[], ClientSession]], ssl_verify: Optional[Union[bool, ssl.SSLContext]] = None, + owns_session: bool = True, ): self.client = client self._ssl_verify = ssl_verify # Store for per-request SSL override - super().__init__(client=client) + super().__init__(client=client, owns_session=owns_session) # Store the client factory for recreating sessions when needed if callable(client): self._client_factory = client diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5cf6efe5ba2..328097639e5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -866,6 +866,7 @@ class AsyncHTTPHandler: return LiteLLMAiohttpTransport( client=shared_session, ssl_verify=ssl_for_transport, + owns_session=False, ) # Create new session only if none provided or existing one is invalid diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 002fa81b9b5..6e2e60ba0dd 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -12,10 +12,42 @@ sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory from litellm.llms.custom_httpx.aiohttp_transport import ( AiohttpResponseStream, + AiohttpTransport, LiteLLMAiohttpTransport, ) +@pytest.mark.asyncio +async def test_aclose_does_not_close_shared_session(): + """Test that aclose() does not close a session it does not own (shared session).""" + session = aiohttp.ClientSession() + try: + transport = LiteLLMAiohttpTransport(client=session, owns_session=False) + await transport.aclose() + assert not session.closed, "Shared session should not be closed by transport" + finally: + await session.close() + + +@pytest.mark.asyncio +async def test_aclose_closes_owned_session(): + """Test that aclose() closes a session it owns.""" + session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=session, owns_session=True) + await transport.aclose() + assert session.closed, "Owned session should be closed by transport" + + +@pytest.mark.asyncio +async def test_owns_session_defaults_to_true(): + """Test that owns_session defaults to True for backwards compatibility.""" + session = aiohttp.ClientSession() + transport = AiohttpTransport(client=session) + assert transport._owns_session is True + await transport.aclose() + assert session.closed + + class MockAiohttpResponse: """Mock aiohttp ClientResponse for testing""" From 06e7bfce2e9fafc17bce023485daa4243bb28bd5 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 14 Feb 2026 11:11:08 -0600 Subject: [PATCH 44/50] perf(spend): avoid duplicate daily agent transaction computation (#21187) --- litellm/proxy/db/db_spend_update_writer.py | 7 ---- .../proxy/db/test_db_spend_update_writer.py | 41 ++++++++++++++++++- 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index dc928921425..9675b82b145 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1725,13 +1725,6 @@ class DBSpendUpdateWriter: "prisma_client is None. Skipping writing spend logs to db." ) return - base_daily_transaction = ( - await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "agent" - ) - ) - if base_daily_transaction is None: - return if payload["agent_id"] is None: verbose_proxy_logger.debug( "agent_id is None for request. Skipping incrementing agent spend." diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 6ccecf59eed..1dd5cba2c4b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -756,6 +756,45 @@ async def test_add_spend_log_transaction_to_daily_agent_transaction_injects_agen assert transaction["custom_llm_provider"] == "openai" +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_agent_transaction_calls_common_helper_once(): + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + payload = { + "request_id": "req-common-helper", + "agent_id": "agent-abc", + "user": "test-user", + "startTime": "2024-01-01T12:00:00", + "api_key": "test-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "model_group": "gpt-4-group", + "prompt_tokens": 12, + "completion_tokens": 6, + "spend": 0.25, + "metadata": '{"usage_object": {}}', + } + + writer.daily_agent_spend_update_queue.add_update = AsyncMock() + original_common_helper = ( + writer._common_add_spend_log_transaction_to_daily_transaction + ) + writer._common_add_spend_log_transaction_to_daily_transaction = AsyncMock( + wraps=original_common_helper + ) + + await writer.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload, + prisma_client=mock_prisma, + ) + + assert ( + writer._common_add_spend_log_transaction_to_daily_transaction.await_count == 1 + ) + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_agent_transaction_skips_when_agent_id_missing(): """ @@ -960,4 +999,4 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): entity_id_field="user_id", table_name="litellm_dailyuserspend", unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", - ) \ No newline at end of file + ) From ac2f17a51413ab640d024510289bb933b6a741ed Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 18:41:06 +0530 Subject: [PATCH 45/50] fix: proxy/batches_endpoints/endpoints.py:309:11: PLR0915 Too many statements (54 > 50) --- litellm/proxy/batches_endpoints/endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 783ac9d6f19..143b2607feb 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -306,7 +306,7 @@ async def create_batch( # noqa: PLR0915 dependencies=[Depends(user_api_key_auth)], tags=["batch"], ) -async def retrieve_batch( +async def retrieve_batch( # noqa: PLR0915 request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), From 4548d9fbe982745908aa368efb172e8f45d93f0d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 18:55:21 +0530 Subject: [PATCH 46/50] fix mypy --- litellm/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 1399bfbeb75..dae0bb1c2c0 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1896,7 +1896,7 @@ def batch_cost_calculator( usage: Usage, model: str, custom_llm_provider: Optional[str] = None, - model_info: Optional[dict] = None, + model_info: Optional[ModelInfo] = None, ) -> Tuple[float, float]: """ Calculate the cost of a batch job. From a38b4c892427dfa2f207383d1f036f582f9a8f9a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 18:59:39 +0530 Subject: [PATCH 47/50] Add doc for OpenAI Agents SDK with LiteLLM --- .../my-website/docs/projects/openai-agents.md | 115 ++++++++++++++++-- docs/my-website/sidebars.js | 1 + 2 files changed, 108 insertions(+), 8 deletions(-) diff --git a/docs/my-website/docs/projects/openai-agents.md b/docs/my-website/docs/projects/openai-agents.md index 95a2191b883..960e8a77551 100644 --- a/docs/my-website/docs/projects/openai-agents.md +++ b/docs/my-website/docs/projects/openai-agents.md @@ -1,22 +1,121 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; -# OpenAI Agents SDK +# OpenAI Agents SDK with LiteLLM -The [OpenAI Agents SDK](https://github.com/openai/openai-agents-python) is a lightweight framework for building multi-agent workflows. -It includes an official LiteLLM extension that lets you use any of the 100+ supported providers (Anthropic, Gemini, Mistral, Bedrock, etc.) +Use OpenAI Agents SDK with any LLM provider through LiteLLM Proxy. + +The [OpenAI Agents SDK](https://github.com/openai/openai-agents-python) is a lightweight framework for building multi-agent workflows. It includes an official LiteLLM extension that lets you use any of the 100+ supported providers. + +## Quick Start + +### 1. Install Dependencies + +```bash +pip install "openai-agents[litellm]" +``` + +### 2. Add Model to Config + +```yaml title="config.yaml" +model_list: + - model_name: gpt-4o + litellm_params: + model: "openai/gpt-4o" + api_key: "os.environ/OPENAI_API_KEY" + + - model_name: claude-sonnet + litellm_params: + model: "anthropic/claude-3-5-sonnet-20241022" + api_key: "os.environ/ANTHROPIC_API_KEY" + + - model_name: gemini-pro + litellm_params: + model: "gemini/gemini-2.0-flash-exp" + api_key: "os.environ/GEMINI_API_KEY" +``` + +### 3. Start LiteLLM Proxy + +```bash +litellm --config config.yaml +``` + +### 4. Use with Proxy + + + ```python from agents import Agent, Runner from agents.extensions.models.litellm_model import LitellmModel +# Point to LiteLLM proxy agent = Agent( name="Assistant", instructions="You are a helpful assistant.", - model=LitellmModel(model="provider/model-name") + model=LitellmModel( + model="claude-sonnet", # Model from config.yaml + api_key="sk-1234", # LiteLLM API key + base_url="http://localhost:4000" + ) ) -result = Runner.run_sync(agent, "your_prompt_here") -print("Result:", result.final_output) +result = await Runner.run(agent, "What is LiteLLM?") +print(result.final_output) ``` -- [GitHub](https://github.com/openai/openai-agents-python) -- [LiteLLM Extension Docs](https://openai.github.io/openai-agents-python/ref/extensions/litellm/) + + + +```python +from agents import Agent, Runner +from agents.extensions.models.litellm_model import LitellmModel + +# Use any provider directly +agent = Agent( + name="Assistant", + instructions="You are a helpful assistant.", + model=LitellmModel( + model="anthropic/claude-3-5-sonnet-20241022", + api_key="your-anthropic-key" + ) +) + +result = await Runner.run(agent, "What is LiteLLM?") +print(result.final_output) +``` + + + + +## Track Usage + +Enable usage tracking to monitor token consumption: + +```python +from agents import Agent, ModelSettings +from agents.extensions.models.litellm_model import LitellmModel + +agent = Agent( + name="Assistant", + model=LitellmModel(model="claude-sonnet", api_key="sk-1234"), + model_settings=ModelSettings(include_usage=True) +) + +result = await Runner.run(agent, "Hello") +print(result.context_wrapper.usage) # Token counts +``` + +## Environment Variables + +| Variable | Value | Description | +|----------|-------|-------------| +| `LITELLM_BASE_URL` | `http://localhost:4000` | LiteLLM proxy URL | +| `LITELLM_API_KEY` | `sk-1234` | Your LiteLLM API key | + +## Related Resources + +- [OpenAI Agents SDK Documentation](https://openai.github.io/openai-agents-python/) +- [LiteLLM Extension Docs](https://openai.github.io/openai-agents-python/models/litellm/) +- [LiteLLM Proxy Quick Start](../proxy/quick_start) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 4efb2475755..5e82fa793b5 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -176,6 +176,7 @@ const sidebars = { "tutorials/copilotkit_sdk", "tutorials/google_adk", "tutorials/livekit_xai_realtime", + "litellm/docs/my-website/docs/projects/openai-agents.md" ] }, From 5b6e232da631dea63dc823aa90fd3ea4d549932b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 19:00:53 +0530 Subject: [PATCH 48/50] Add doc for OpenAI Agents SDK with LiteLLM --- docs/my-website/docs/projects/openai-agents.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/docs/projects/openai-agents.md b/docs/my-website/docs/projects/openai-agents.md index 960e8a77551..86983e7e510 100644 --- a/docs/my-website/docs/projects/openai-agents.md +++ b/docs/my-website/docs/projects/openai-agents.md @@ -1,7 +1,7 @@ import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -# OpenAI Agents SDK with LiteLLM +# OpenAI Agents SDK Use OpenAI Agents SDK with any LLM provider through LiteLLM Proxy. From be9df253dd4538c060049188f2ad0cc6de5ac4ad Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 19:03:15 +0530 Subject: [PATCH 49/50] Update docs/my-website/sidebars.js Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- docs/my-website/sidebars.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 5e82fa793b5..42996d1a3e9 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -176,7 +176,7 @@ const sidebars = { "tutorials/copilotkit_sdk", "tutorials/google_adk", "tutorials/livekit_xai_realtime", - "litellm/docs/my-website/docs/projects/openai-agents.md" + "projects/openai-agents" ] }, From 6ceebdcfe28c8a09d59db91d6599af6c3e30460e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 16 Feb 2026 19:05:06 +0530 Subject: [PATCH 50/50] fix mypy --- litellm/batches/batch_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index c92ab9b230e..29bd99c2a60 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -8,7 +8,7 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.types.llms.openai import Batch -from litellm.types.utils import CallTypes, ModelResponse, Usage +from litellm.types.utils import CallTypes, ModelInfo, ModelResponse, Usage from litellm.utils import token_counter @@ -16,7 +16,7 @@ async def calculate_batch_cost_and_usage( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: Optional[str] = None, - model_info: Optional[dict] = None, + model_info: Optional[ModelInfo] = None, ) -> Tuple[float, Usage, List[str]]: """ Calculate the cost and usage of a batch. @@ -102,7 +102,7 @@ def _batch_cost_calculator( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", model_name: Optional[str] = None, - model_info: Optional[dict] = None, + model_info: Optional[ModelInfo] = None, ) -> float: """ Calculate the cost of a batch based on the output file id @@ -300,7 +300,7 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", - model_info: Optional[dict] = None, + model_info: Optional[ModelInfo] = None, ) -> float: """ Get the cost of a batch job from the file content