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).
This commit is contained in:
ryan-crabbe-berri 2026-07-29 19:39:16 -07:00
parent b7ac3e237b
commit 11014897a7
6 changed files with 320 additions and 145 deletions

View file

@ -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":

View file

@ -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):

View file

@ -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(

View file

@ -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)

View file

@ -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": "<VOLATILE>",
}
# 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)

View file

@ -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"""