From 11014897a7b098d0e8d46aefa17d850dded0e19a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 29 Jul 2026 19:39:16 -0700 Subject: [PATCH] fix(proxy): scope reload persistence to the model cost map and seed the pod clock from the actual load time Revert the Anthropic beta headers reload path to its previous JSON-flag implementation so this PR only changes the price data reload; the beta headers path keeps working exactly as before and can migrate to the shared module in a follow-up. The unused columns on its config row are inert. Seed model_cost_map_loaded_at from the timestamp get_model_cost_map records at the actual import-time fetch instead of ProxyConfig construction time, closing the startup window where a manual reload request stamped between the fetch and the constructor compared as older than the pod's data and was skipped (Greptile P1 on the previous commit). --- .../litellm_core_utils/get_model_cost_map.py | 8 + .../common_utils/periodic_reload_schedule.py | 5 +- litellm/proxy/proxy_server.py | 240 ++++++++++++++---- .../test_get_model_cost_map.py | 22 ++ .../test_routes_anthropic_beta.py | 71 ++---- tests/test_litellm/proxy/test_proxy_server.py | 119 +++++---- 6 files changed, 320 insertions(+), 145 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 4c0a01ad645..c9863e78b0e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -10,6 +10,7 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True import json import os +from datetime import datetime, timezone from importlib.resources import files from typing import Dict, List, Optional @@ -169,6 +170,7 @@ class ModelCostMapSourceInfo: url: Optional[str] = None is_env_forced: bool = False fallback_reason: Optional[str] = None + loaded_at: "datetime | None" = None # Module-level singleton tracking the source of the current cost map @@ -193,6 +195,11 @@ def get_model_cost_map_source_info() -> dict: } +def get_model_cost_map_loaded_at() -> "datetime | None": + """When this process last loaded its cost map, stamped at the start of every load""" + return _cost_map_source_info.loaded_at + + def _expand_model_aliases(model_cost: dict) -> dict: """ Expand ``aliases`` lists in model cost entries into top-level entries. @@ -272,6 +279,7 @@ def get_model_cost_map(url: str) -> dict: The full backup dict is only parsed when it must be *returned* as a fallback — it is never held in memory long-term. """ + _cost_map_source_info.loaded_at = datetime.now(timezone.utc) # Note: can't use get_secret_bool here — this runs during litellm.__init__ # before litellm._key_management_settings is set. if os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": diff --git a/litellm/proxy/common_utils/periodic_reload_schedule.py b/litellm/proxy/common_utils/periodic_reload_schedule.py index 44db004bd4f..26ae6bbc4a7 100644 --- a/litellm/proxy/common_utils/periodic_reload_schedule.py +++ b/litellm/proxy/common_utils/periodic_reload_schedule.py @@ -1,6 +1,6 @@ """ -Persistence for the admin-configured periodic reload schedules (model cost map, -Anthropic beta headers) stored in ``LiteLLM_Config``. +Persistence for the admin-configured periodic model cost map reload schedule stored in +``LiteLLM_Config``. Field ownership is split by writer so concurrent writers never overwrite each other: the schedule endpoints own the ``param_value`` JSON (``interval_hours``), while the @@ -34,7 +34,6 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_Config MODEL_COST_MAP_RELOAD_PARAM_NAME = "model_cost_map_reload_config" -ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME = "anthropic_beta_headers_reload_config" class _ConfigRowWrite(TypedDict, total=False): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77123003d08..a28f39c7a86 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -327,7 +327,6 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) from litellm.proxy.common_utils.periodic_reload_schedule import ( - ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, MODEL_COST_MAP_RELOAD_PARAM_NAME, pod_reload_is_due, read_reload_schedule, @@ -2067,6 +2066,9 @@ celery_fn = None # Redis Queue for handling requests scheduler = None +# Global variable for anthropic beta headers reload scheduling +last_anthropic_beta_headers_reload = None + ### DB WRITER ### db_writer_client: Optional[AsyncHTTPHandler] = None @@ -3856,8 +3858,11 @@ class ProxyConfig: self._last_semantic_filter_config: Optional[Dict[str, Any]] = None self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None self.worker_registry: List["WorkerRegistryEntry"] = [] - self.model_cost_map_loaded_at: datetime = utc_now() - self.anthropic_beta_headers_loaded_at: datetime = utc_now() + from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map_loaded_at, + ) + + self.model_cost_map_loaded_at: datetime = get_model_cost_map_loaded_at() or utc_now() def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -6197,6 +6202,7 @@ class ProxyConfig: "router_settings", "litellm_settings", "environment_variables", + "anthropic_beta_headers_reload_config", ], ) @@ -6259,6 +6265,9 @@ class ProxyConfig: if self._should_load_db_object(object_type="tools"): await self._init_tool_policy_in_db(prisma_client=prisma_client) + if self._should_load_db_object(object_type="anthropic_beta_headers"): + await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) + if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="cache_settings"): @@ -6419,17 +6428,14 @@ class ProxyConfig: async def check_periodic_reloads(self, prisma_client: PrismaClient): """ - Run the admin-configured periodic reloads (model cost map, Anthropic beta headers). + Run the admin-configured periodic model cost map reload. - Scheduled on its own job so schedules configured from the Admin UI fire whether + Scheduled on its own job so a schedule configured from the Admin UI fires whether or not `store_model_in_db` is enabled. """ if self._should_load_db_object(object_type="model_cost_map"): await self._check_and_reload_model_cost_map(prisma_client=prisma_client) - if self._should_load_db_object(object_type="anthropic_beta_headers"): - await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) - async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): """ Check if model cost map needs to be reloaded based on database configuration. @@ -6462,35 +6468,96 @@ class ProxyConfig: async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient): """ Check if anthropic beta headers config needs to be reloaded based on database configuration. - Runs on the periodic reload job, independently of `store_model_in_db`. + This function runs every 10 seconds as part of _init_non_llm_objects_in_db. """ try: - schedule = await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) - if schedule is None: - return + # Get anthropic beta headers reload configuration from database + config_record = await get_config_param(prisma_client, "anthropic_beta_headers_reload_config") - current_time = utc_now() - if not pod_reload_is_due( - schedule=schedule, - pod_data_loaded_at=self.anthropic_beta_headers_loaded_at, - current_time=current_time, - description="Anthropic beta headers", - ): - return + if config_record is None or config_record.param_value is None: + return # No configuration found, skip reload - from litellm.anthropic_beta_headers_manager import ( - reload_beta_headers_config, - ) + config = config_record.param_value + interval_hours = config.get("interval_hours") + force_reload = config.get("force_reload", False) - new_config = await asyncio.to_thread(reload_beta_headers_config) + if interval_hours is None and force_reload is False: + return # No interval configured, skip reload - self.anthropic_beta_headers_loaded_at = current_time - await record_reload_run(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time) + current_time = datetime.utcnow() - provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") - verbose_proxy_logger.info( - f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}" - ) + # Check if we need to reload based on interval or force reload + should_reload = False + + if force_reload: + should_reload = True + verbose_proxy_logger.info("Anthropic beta headers reload triggered by force reload flag") + elif interval_hours is not None: + # Use pod's in-memory last reload time + global last_anthropic_beta_headers_reload + if last_anthropic_beta_headers_reload is not None: + try: + last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) + time_since_last_reload = current_time - last_reload_time + hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 + + if hours_since_last_reload >= interval_hours: + should_reload = True + verbose_proxy_logger.info( + f"Anthropic beta headers reload triggered by interval. Hours since last reload: {hours_since_last_reload:.2f}, Interval: {interval_hours}" + ) + except Exception as e: + verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") + # If we can't parse the last reload time, reload anyway + should_reload = True + else: + # No last reload time recorded, reload now + should_reload = True + verbose_proxy_logger.info( + "Anthropic beta headers reload triggered - no previous reload time recorded" + ) + + if should_reload: + # Perform the reload + from litellm.anthropic_beta_headers_manager import ( + reload_beta_headers_config, + ) + + new_config = reload_beta_headers_config() + + # Update pod's in-memory last reload time + last_anthropic_beta_headers_reload = current_time.isoformat() + + # Clear force reload flag in database + await ConfigRepository(prisma_client).table.upsert( + where={"param_name": "anthropic_beta_headers_reload_config"}, + data={ + "create": { + "param_name": "anthropic_beta_headers_reload_config", + "param_value": safe_dumps( + { + "interval_hours": interval_hours, + "force_reload": False, + } + ), + }, + "update": { + "param_value": safe_dumps( + { + "interval_hours": interval_hours, + "force_reload": False, + } + ) + }, + }, + ) + await invalidate_config_param("anthropic_beta_headers_reload_config") + + # Count providers in config + provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") + verbose_proxy_logger.info( + f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}" + ) except Exception as e: verbose_proxy_logger.exception(f"Error in _check_and_reload_anthropic_beta_headers: {str(e)}") @@ -15797,13 +15864,32 @@ async def reload_anthropic_beta_headers( # Immediately reload the beta headers config in the current pod from litellm.anthropic_beta_headers_manager import reload_beta_headers_config - new_config = await asyncio.to_thread(reload_beta_headers_config) + new_config = reload_beta_headers_config() - current_time = utc_now() - proxy_config.anthropic_beta_headers_loaded_at = current_time + # Update pod's in-memory last reload time + global last_anthropic_beta_headers_reload + current_time = datetime.utcnow() + last_anthropic_beta_headers_reload = current_time.isoformat() - # Stamp the shared last run and the reload request so other pods reload too - await record_manual_reload(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, current_time) + # Set force reload flag in database for other pods, preserving existing interval_hours + existing_beta_config = await ConfigRepository(prisma_client).table.find_unique( + where={"param_name": "anthropic_beta_headers_reload_config"} + ) + existing_beta_interval = None + if existing_beta_config and existing_beta_config.param_value: + existing_beta_interval = existing_beta_config.param_value.get("interval_hours") + + await ConfigRepository(prisma_client).table.upsert( + where={"param_name": "anthropic_beta_headers_reload_config"}, + data={ + "create": { + "param_name": "anthropic_beta_headers_reload_config", + "param_value": safe_dumps({"interval_hours": None, "force_reload": True}), + }, + "update": {"param_value": safe_dumps({"interval_hours": existing_beta_interval, "force_reload": True})}, + }, + ) + await invalidate_config_param("anthropic_beta_headers_reload_config") provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"]) verbose_proxy_logger.info( @@ -15852,7 +15938,18 @@ async def schedule_anthropic_beta_headers_reload( if prisma_client is None: raise HTTPException(status_code=500, detail="Database connection not available") - await write_reload_interval(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME, hours) + # Update database with new reload configuration + await ConfigRepository(prisma_client).table.upsert( + where={"param_name": "anthropic_beta_headers_reload_config"}, + data={ + "create": { + "param_name": "anthropic_beta_headers_reload_config", + "param_value": safe_dumps({"interval_hours": hours, "force_reload": False}), + }, + "update": {"param_value": safe_dumps({"interval_hours": hours, "force_reload": False})}, + }, + ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info(f"Anthropic beta headers reload scheduled for every {hours} hours") @@ -15860,7 +15957,7 @@ async def schedule_anthropic_beta_headers_reload( "message": f"Anthropic beta headers reload scheduled for every {hours} hours", "status": "success", "interval_hours": hours, - "timestamp": utc_now().isoformat(), + "timestamp": datetime.utcnow().isoformat(), } except Exception as e: verbose_proxy_logger.exception(f"Failed to schedule anthropic beta headers reload: {str(e)}") @@ -15897,17 +15994,15 @@ async def cancel_anthropic_beta_headers_reload( raise HTTPException(status_code=500, detail="Database connection not available") # Remove reload configuration from database - await ConfigRepository(prisma_client).table.delete( - where={"param_name": ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME} - ) - await invalidate_config_param(ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) + await ConfigRepository(prisma_client).table.delete(where={"param_name": "anthropic_beta_headers_reload_config"}) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled") return { "message": "Anthropic beta headers reload schedule cancelled", "status": "success", - "timestamp": utc_now().isoformat(), + "timestamp": datetime.utcnow().isoformat(), } except Exception as e: verbose_proxy_logger.exception(f"Failed to cancel anthropic beta headers reload: {str(e)}") @@ -15939,15 +16034,68 @@ async def get_anthropic_beta_headers_reload_status( ) try: - global prisma_client + global prisma_client, last_anthropic_beta_headers_reload + + verbose_proxy_logger.info( + f"Checking anthropic beta headers reload status. Last reload: {last_anthropic_beta_headers_reload}" + ) if prisma_client is None: verbose_proxy_logger.info("No database connection, returning not scheduled") - return reload_schedule_status(None) + return { + "scheduled": False, + "interval_hours": None, + "last_run": None, + "next_run": None, + } - return reload_schedule_status( - await read_reload_schedule(prisma_client, ANTHROPIC_BETA_HEADERS_RELOAD_PARAM_NAME) + # Get reload configuration from database + config_record = await ConfigRepository(prisma_client).table.find_unique( + where={"param_name": "anthropic_beta_headers_reload_config"} ) + + if config_record is None or config_record.param_value is None: + verbose_proxy_logger.info("No anthropic beta headers reload configuration found") + return { + "scheduled": False, + "interval_hours": None, + "last_run": None, + "next_run": None, + } + + config = config_record.param_value + interval_hours = config.get("interval_hours") + + if interval_hours is None: + verbose_proxy_logger.info("No interval configured, returning not scheduled") + return { + "scheduled": False, + "interval_hours": None, + "last_run": None, + "next_run": None, + } + + current_time = datetime.utcnow() + next_run = None + + # Use pod's in-memory last reload time + if last_anthropic_beta_headers_reload is not None: + try: + last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) + time_since_last_reload = current_time - last_reload_time + hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 + + if hours_since_last_reload < interval_hours: + next_run = (last_reload_time + timedelta(hours=interval_hours)).isoformat() + except Exception as e: + verbose_proxy_logger.warning(f"Error parsing last reload time: {e}") + + return { + "scheduled": True, + "interval_hours": interval_hours, + "last_run": last_anthropic_beta_headers_reload, + "next_run": next_run, + } except Exception as e: verbose_proxy_logger.exception(f"Failed to get anthropic beta headers reload status: {str(e)}") raise HTTPException( diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index f21c667276e..f2729bdacd8 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -248,3 +248,25 @@ def test_azure_ai_claude_1m_context_entries(cost_map: dict): "azure_ai/claude-haiku-4-5", ]: assert cost_map[model]["max_input_tokens"] == 200000, model + + +def test_get_model_cost_map_stamps_loaded_at(monkeypatch): + """The load time feeds each pod's reload-due decision; a load that does not stamp it + would make manual reload requests race the proxy's startup""" + from datetime import datetime, timezone + + from litellm.litellm_core_utils import get_model_cost_map as module + + monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None) + monkeypatch.setattr( + module.GetModelCostMap, + "fetch_remote_model_cost_map", + staticmethod(lambda url, timeout=5: _load_root_cost_map()), + ) + + before = datetime.now(timezone.utc) + module.get_model_cost_map(url="https://example.invalid/cost_map.json") + loaded_at = module.get_model_cost_map_loaded_at() + + assert loaded_at is not None + assert before <= loaded_at <= datetime.now(timezone.utc) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py index 2e729751d66..7ef29b71bf0 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_anthropic_beta.py @@ -9,9 +9,11 @@ Routes covered: from __future__ import annotations +import json from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock +import pytest from .conftest import VOLATILE_KEYS, normalize @@ -97,24 +99,22 @@ def test_reload_anthropic_beta_headers_admin_success(client, auth_as, monkeypatc "providers_count": 2, "timestamp": "", } + # And the upsert was actually invoked (force_reload write). prisma.db.litellm_config.upsert.assert_awaited_once() def test_reload_anthropic_beta_headers_preserves_existing_interval( client, auth_as, monkeypatch ): - """The manual-reload write only stamps the job-owned timestamp columns; param_value - (the admin-owned interval) is not in the payload, so it cannot be clobbered.""" - from datetime import datetime - + """When an existing reload config has an interval set, the force-reload + write must preserve that interval (the route reads it back then upserts + with the same number). This pins the read-then-write behaviour.""" from litellm.proxy._types import LitellmUserRoles _stub_reload_beta_headers(monkeypatch) existing = SimpleNamespace( param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": 12}, - reload_requested_at=None, - last_run_at=None, + param_value={"interval_hours": 12, "force_reload": False}, ) prisma = _make_prisma_with_config(config_record=existing) _install_prisma(monkeypatch, prisma) @@ -123,12 +123,17 @@ def test_reload_anthropic_beta_headers_preserves_existing_interval( response = client.post("/reload/anthropic_beta_headers") assert response.status_code == 200 + # The update branch's interval_hours was sourced from the existing record. call_kwargs = prisma.db.litellm_config.upsert.await_args.kwargs - update_payload = call_kwargs["data"]["update"] - assert set(update_payload) == {"last_run_at", "reload_requested_at"} - assert isinstance(update_payload["last_run_at"], datetime) - assert update_payload["reload_requested_at"] == update_payload["last_run_at"] - assert "param_value" not in call_kwargs["data"]["create"] + data = call_kwargs["data"] + update_payload = data["update"]["param_value"] + parsed = ( + json.loads(update_payload) + if isinstance(update_payload, str) + else update_payload + ) + assert parsed["interval_hours"] == 12 + assert parsed["force_reload"] is True def test_reload_anthropic_beta_headers_not_admin_forbidden(client, auth_as): @@ -286,16 +291,17 @@ def test_get_anthropic_beta_headers_reload_status_scheduled( ): """When a config row with ``interval_hours`` is present, ``scheduled`` is True and ``interval_hours`` echoes the DB value. Pins the full response shape.""" + from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles record = SimpleNamespace( param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": 6}, - reload_requested_at=None, - last_run_at=None, + param_value={"interval_hours": 6, "force_reload": False}, ) prisma = _make_prisma_with_config(config_record=record) _install_prisma(monkeypatch, prisma) + # No prior reload — next_run stays None. + monkeypatch.setattr(ps, "last_anthropic_beta_headers_reload", None) with auth_as(LitellmUserRoles.PROXY_ADMIN): response = client.get("/schedule/anthropic_beta_headers_reload/status") @@ -309,35 +315,6 @@ def test_get_anthropic_beta_headers_reload_status_scheduled( } -def test_get_anthropic_beta_headers_reload_status_reports_persisted_last_run( - client, auth_as, monkeypatch -): - """last_run/next_run come from the DB row, so status survives pod restarts.""" - from datetime import datetime, timezone - - from litellm.proxy._types import LitellmUserRoles - - record = SimpleNamespace( - param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": 6}, - reload_requested_at=None, - last_run_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc), - ) - prisma = _make_prisma_with_config(config_record=record) - _install_prisma(monkeypatch, prisma) - - with auth_as(LitellmUserRoles.PROXY_ADMIN): - response = client.get("/schedule/anthropic_beta_headers_reload/status") - - assert response.status_code == 200 - assert normalize(response.json()) == { - "scheduled": True, - "interval_hours": 6, - "last_run": "2024-01-01T06:00:00+00:00", - "next_run": "2024-01-01T12:00:00+00:00", - } - - def test_get_anthropic_beta_headers_reload_status_not_scheduled_no_db( client, auth_as, monkeypatch ): @@ -362,15 +339,11 @@ def test_get_anthropic_beta_headers_reload_status_no_interval_unscheduled( client, auth_as, monkeypatch ): """Config row present but ``interval_hours`` is None → unscheduled response.""" - from datetime import datetime, timezone - from litellm.proxy._types import LitellmUserRoles record = SimpleNamespace( param_name="anthropic_beta_headers_reload_config", - param_value={"interval_hours": None}, - reload_requested_at=datetime(2024, 1, 1, 6, 0, 0, tzinfo=timezone.utc), - last_run_at=None, + param_value={"interval_hours": None, "force_reload": True}, ) prisma = _make_prisma_with_config(config_record=record) _install_prisma(monkeypatch, prisma) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index dc7ddc0a464..552cab81c42 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3972,6 +3972,20 @@ class TestPriceDataReloadIntegration: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() + def test_pod_data_clock_seeded_from_actual_cost_map_load(self): + """Regression: seeding from ProxyConfig construction time instead of the real + import-time fetch let a manual request stamped during startup be skipped""" + from datetime import datetime, timezone + + from litellm.proxy.proxy_server import ProxyConfig + + fetch_time = datetime(2024, 1, 1, 6, 0, tzinfo=timezone.utc) + with patch( + "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map_loaded_at", + return_value=fetch_time, + ): + assert ProxyConfig().model_cost_map_loaded_at == fetch_time + def test_distributed_reload_check_function(self): """ A reload requested by another pod after this pod's last reload takes effect here @@ -4215,39 +4229,52 @@ class TestPriceDataReloadIntegration: _invalidate_model_cost_lowercase_map() def test_anthropic_beta_headers_reload_preserves_interval_hours(self): - """Same column-only write contract as the cost map job, for the beta headers schedule""" - from litellm.proxy.proxy_server import ProxyConfig + """Test that _check_and_reload_anthropic_beta_headers preserves interval_hours after reload. + Regression test: the update branch of the upsert was dropping interval_hours, + identical to the model cost map bug. + """ + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import litellm_config_cache + + litellm_config_cache.flush_cache() proxy_config = ProxyConfig() mock_prisma = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=_reload_schedule_row({"interval_hours": 12}) - ) - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) - frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) - with ( - patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload, - patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), - ): + # Set up config with interval_hours=12 and force_reload=True to trigger reload + mock_config = MagicMock() + mock_config.param_value = {"interval_hours": 12, "force_reload": True} + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_anthropic_beta_headers now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) + + with patch( + "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" + ) as mock_reload: mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} asyncio.run( proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma) ) - assert proxy_config.anthropic_beta_headers_loaded_at == frozen_now - assert mock_prisma.db.litellm_config.update_many.call_args[1] == { - "data": {"last_run_at": frozen_now}, - "where": {"param_name": "anthropic_beta_headers_reload_config"}, - } - mock_prisma.db.litellm_config.upsert.assert_not_called() + # Verify the upsert update branch preserves interval_hours + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == False + assert param_value_dict["interval_hours"] == 12, ( + "interval_hours must be preserved in the update branch; " + "dropping it causes the schedule to self-destruct" + ) def test_anthropic_beta_headers_manual_reload_preserves_interval_hours(self): - """Manual beta headers reload stamps the run columns only, leaving the schedule intact""" + """Test that manual reload via /reload/anthropic_beta_headers preserves existing interval_hours. + + Regression test: the manual reload endpoint was overwriting param_value with + only force_reload=True, dropping any existing interval_hours schedule. + """ from litellm.proxy._types import LitellmUserRoles from litellm.proxy.proxy_server import cleanup_router_config_variables @@ -4260,36 +4287,34 @@ class TestPriceDataReloadIntegration: mock_auth.user_role = LitellmUserRoles.PROXY_ADMIN app.dependency_overrides[user_api_key_auth] = lambda: mock_auth client = TestClient(app) - frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) - with ( - patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload, - patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, - patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), - ): + with patch( + "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" + ) as mock_reload: mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} - mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - response = client.post("/reload/anthropic_beta_headers") - assert response.status_code == 200 + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + # Simulate existing config with a schedule + mock_existing = MagicMock() + mock_existing.param_value = {"interval_hours": 8, "force_reload": False} + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=mock_existing + ) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) - mock_prisma.db.litellm_config.find_unique.assert_not_called() - call_args = mock_prisma.db.litellm_config.upsert.call_args - assert call_args[1]["data"]["update"] == { - "last_run_at": frozen_now, - "reload_requested_at": frozen_now, - } - assert call_args[1]["data"]["create"] == { - "param_name": "anthropic_beta_headers_reload_config", - "last_run_at": frozen_now, - "reload_requested_at": frozen_now, - } - assert ( - proxy_server_module.proxy_config.anthropic_beta_headers_loaded_at - == frozen_now - ) + response = client.post("/reload/anthropic_beta_headers") + assert response.status_code == 200 + + # Verify interval_hours was preserved in the upsert + mock_prisma.db.litellm_config.upsert.assert_called() + call_args = mock_prisma.db.litellm_config.upsert.call_args + param_value_json = call_args[1]["data"]["update"]["param_value"] + param_value_dict = json.loads(param_value_json) + assert param_value_dict["force_reload"] == True + assert param_value_dict["interval_hours"] == 8, ( + "interval_hours must be preserved when manual reload sets force_reload; " + "dropping it destroys any existing schedule" + ) def test_config_file_parsing(self): """Test parsing of config file with reload settings"""