From 194c984cff053fc35235d7491a6ad5c59e685564 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 12 Nov 2026 00:54:10 +0000 Subject: [PATCH] fix(proxy): keep a cost settings update from being undone by an in-flight config reload (#51273) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cost_tracking_settings.py | 61 ++++++++----------- .../test_cost_tracking_settings.py | 56 ++++++++++++++++- 2 files changed, 79 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 7e6aec84292..dcc6a614a54 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -253,6 +253,7 @@ async def update_cost_discount_config( ``` """ from litellm.proxy.proxy_server import ( + DB_RECONCILE_LOCK, prisma_client, proxy_config, store_model_in_db, @@ -301,21 +302,13 @@ async def update_cost_discount_config( ) try: - # Load existing config - config: Final = await proxy_config.get_config() - - # Ensure litellm_settings exists - if "litellm_settings" not in config: - config["litellm_settings"] = {} - - # Update cost_discount_config - config["litellm_settings"]["cost_discount_config"] = cost_discount_config - - # Save the updated config to DB - await proxy_config.save_config(new_config=config) - - # Update in-memory litellm.cost_discount_config - litellm.cost_discount_config = cost_discount_config + async with DB_RECONCILE_LOCK: + config: Final = await proxy_config.get_config() + if "litellm_settings" not in config: + config["litellm_settings"] = {} + config["litellm_settings"]["cost_discount_config"] = cost_discount_config + await proxy_config.save_config(new_config=config) + litellm.cost_discount_config = cost_discount_config verbose_proxy_logger.info("Updated cost_discount_config: %s", cost_discount_config) @@ -397,6 +390,7 @@ async def update_cost_margin_config( ``` """ from litellm.proxy.proxy_server import ( + DB_RECONCILE_LOCK, prisma_client, proxy_config, store_model_in_db, @@ -475,21 +469,13 @@ async def update_cost_margin_config( ) try: - # Load existing config - config: Final = await proxy_config.get_config() - - # Ensure litellm_settings exists - if "litellm_settings" not in config: - config["litellm_settings"] = {} - - # Update cost_margin_config - config["litellm_settings"]["cost_margin_config"] = cost_margin_config - - # Save the updated config to DB - await proxy_config.save_config(new_config=config) - - # Update in-memory litellm.cost_margin_config - litellm.cost_margin_config = cost_margin_config + async with DB_RECONCILE_LOCK: + config: Final = await proxy_config.get_config() + if "litellm_settings" not in config: + config["litellm_settings"] = {} + config["litellm_settings"]["cost_margin_config"] = cost_margin_config + await proxy_config.save_config(new_config=config) + litellm.cost_margin_config = cost_margin_config verbose_proxy_logger.info("Updated cost_margin_config: %s", cost_margin_config) @@ -534,6 +520,7 @@ async def update_block_requests_for_models_without_pricing( request: BlockUnpricedModelsRequest, ) -> BlockUnpricedModelsResponse: from litellm.proxy.proxy_server import ( + DB_RECONCILE_LOCK, prisma_client, proxy_config, store_model_in_db, @@ -552,13 +539,13 @@ async def update_block_requests_for_models_without_pricing( ) try: - config = await proxy_config.get_config() - if "litellm_settings" not in config: - config["litellm_settings"] = {} - config["litellm_settings"]["block_requests_for_models_without_pricing"] = request.enabled - await proxy_config.save_config(new_config=config) - - litellm.block_requests_for_models_without_pricing = request.enabled + async with DB_RECONCILE_LOCK: + config = await proxy_config.get_config() + if "litellm_settings" not in config: + config["litellm_settings"] = {} + config["litellm_settings"]["block_requests_for_models_without_pricing"] = request.enabled + await proxy_config.save_config(new_config=config) + litellm.block_requests_for_models_without_pricing = request.enabled verbose_proxy_logger.info("Updated block_requests_for_models_without_pricing: %s", request.enabled) return BlockUnpricedModelsResponse(enabled=request.enabled) diff --git a/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py index b93e23dcb94..8b06e6545a1 100644 --- a/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/unit/proxy/management_endpoints/test_cost_tracking_settings.py @@ -4,6 +4,7 @@ Tests for cost tracking settings management endpoints. Tests the GET and PATCH endpoints for managing cost discount configuration. """ +import asyncio from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -15,7 +16,11 @@ from pydantic import ValidationError import litellm from litellm._internal_context import pinned_billing_time from litellm.proxy._types import CostEstimateRequest -from litellm.proxy.management_endpoints.cost_tracking_settings import router +from litellm.proxy.management_endpoints.cost_tracking_settings import ( + router, + update_cost_discount_config, + update_cost_margin_config, +) from litellm.proxy.proxy_server import app client = TestClient(app) @@ -762,6 +767,55 @@ class TestEstimateCostOnPremProvider: +class TestCostSettingsPatchAgainstConfigReload: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("setting", "update", "stale", "fresh"), + [ + ( + "cost_margin_config", + update_cost_margin_config, + {"openai": {"percentage": 0.1, "fixed_amount": 0.001}}, + {"global": 1}, + ), + ("cost_discount_config", update_cost_discount_config, {"openai": 0.05}, {}), + ], + ) + async def test_patch_is_not_undone_by_a_reload_that_read_the_previous_row(self, setting, update, stale, fresh): + from litellm.proxy import proxy_server + + reload_has_read_the_row = asyncio.Event() + release_reload = asyncio.Event() + + async def reload_holding_the_previous_row() -> None: + async with proxy_server.DB_RECONCILE_LOCK: + reload_has_read_the_row.set() + await release_reload.wait() + proxy_server.ProxyConfig()._apply_litellm_settings_db_values({setting: stale}) + + mock_proxy_config = AsyncMock() + mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {setting: stale}}) + mock_proxy_config.save_config = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.DB_RECONCILE_LOCK", asyncio.Lock()), + patch.object(litellm, setting, stale), + ): + reload = asyncio.create_task(reload_holding_the_previous_row()) + await reload_has_read_the_row.wait() + patched = asyncio.create_task(update(fresh, user_api_key_dict=MagicMock())) + for _ in range(20): + await asyncio.sleep(0) + release_reload.set() + await asyncio.gather(reload, patched) + + assert patched.result()["values"] == fresh + assert getattr(litellm, setting) == fresh + + class TestBlockRequestsForModelsWithoutPricing: """Test suite for the block_requests_for_models_without_pricing toggle endpoints"""