fix(proxy): compare the reset against the catalog as loaded, not only the bundled backup

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/litellm_core_utils/get_model_cost_map.py

@@ -10,7 +10,6 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True
 """
 
 import asyncio
-import functools
 import hashlib
 import json
 import os
@@ -40,7 +39,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
 )
 
 FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations"
-_BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]])
+_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]])
 _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"})
 
 
@@ -92,12 +91,17 @@ class GetModelCostMap:
         """Load the local backup model cost map bundled with the package."""
         return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map
 
-    @staticmethod
-    @functools.lru_cache(maxsize=1)
-    def bundled_model_cost_map() -> Mapping[str, Mapping[str, object]]:
-        """The bundled catalog as shipped, untouched by ``register_model`` or router registrations."""
-        raw: Final = _BUNDLED_CATALOG_ADAPTER.validate_python(GetModelCostMap.load_local_model_cost_map())
-        return MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()})
+    _loaded_catalog: Mapping[str, Mapping[str, object]] = MappingProxyType({})
+
+    @classmethod
+    def loaded_model_cost_map(cls) -> Mapping[str, Mapping[str, object]]:
+        """The catalog as last loaded (bundled or remote), untouched by ``register_model`` or router registrations."""
+        return cls._loaded_catalog
+
+    @classmethod
+    def _snapshot_loaded_catalog(cls, model_cost: dict) -> None:
+        raw: Final = _CATALOG_ADAPTER.validate_python(model_cost)
+        cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()})
 
     @classmethod
     def _get_backup_model_count(cls) -> int:
@@ -544,7 +548,9 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
 def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded:
     _cost_map_source_info.source_revision = loaded.revision
     _cost_map_source_info.etag = loaded.etag
-    return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map))
+    finalized: Final = _finalize_model_cost_map(loaded.model_cost_map)
+    GetModelCostMap._snapshot_loaded_catalog(finalized)  # pyright: ignore[reportPrivateUsage]  # same module
+    return replace(loaded, model_cost_map=finalized)
 
 
 def adopt_model_cost_map(

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -888,11 +888,11 @@ def _cost_map_entry(db_model: Deployment, incoming_model_info: Mapping[str, obje
         return MappingProxyType({})
 
 
-def _bundled_cost_map_entry(incoming_model_info: Mapping[str, object]) -> Mapping[str, object]:
+def _loaded_catalog_entry(incoming_model_info: Mapping[str, object]) -> Mapping[str, object]:
     catalog_key: Final = incoming_model_info.get(COST_MAP_LOOKUP_KEY)
     if not isinstance(catalog_key, str):
         return MappingProxyType({})
-    return GetModelCostMap.bundled_model_cost_map().get(catalog_key, MappingProxyType({}))
+    return GetModelCostMap.loaded_model_cost_map().get(catalog_key, MappingProxyType({}))
 
 
 def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel:
@@ -921,7 +921,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
         echoed_fields: Final = echoed_cost_map_fields(
             incoming_model_info,
             _cost_map_entry(db_model, incoming_model_info),
-            _bundled_cost_map_entry(incoming_model_info),
+            _loaded_catalog_entry(incoming_model_info),
         )
         merged_model_info.update(
             MappingProxyType(

litellm/types/utils.py

@@ -3786,7 +3786,7 @@ def echoed_cost_map_fields(model_info: Mapping[str, Any], *cost_map_entries: Map
     response. Anything in it that still equals a resolved cost-map entry is a display value
     nobody typed; a value the operator edited differs from every entry and stays a real override.
     Callers pass both the live entry, which the router rewrites with each deployment's own
-    overrides, and the bundled entry, so a reset to the catalog value reads as an echo either way.
+    overrides, and the catalog entry as loaded, so a reset to the catalog value reads as an echo either way.
     """
     if COST_MAP_LOOKUP_KEY not in model_info:
         return ()

tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py

@@ -448,6 +448,23 @@ async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_eta
     assert get_model_cost_map_provenance() == {"source_revision": git_blob_id(body), "etag": 'W/"abc123"'}
 
 
+@pytest.mark.asyncio
+async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later_registrations(monkeypatch):
+    import litellm
+
+    edited = json.loads(_real_map_bytes())
+    edited["gpt-5.4-mini"]["max_input_tokens"] = 777
+    client, _ = _mock_client([httpx.Response(200, content=json.dumps(edited).encode())])
+
+    result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client)
+
+    assert isinstance(result, ModelCostMapReloaded)
+    monkeypatch.setattr(litellm, "model_cost", result.model_cost_map)
+    litellm.register_model({"gpt-5.4-mini": {"max_input_tokens": 2048}}, persist_across_reloads=False)
+    assert litellm.model_cost["gpt-5.4-mini"]["max_input_tokens"] == 2048
+    assert GetModelCostMap.loaded_model_cost_map()["gpt-5.4-mini"]["max_input_tokens"] == 777
+
+
 @pytest.mark.asyncio
 async def test_refetch_revision_follows_the_bytes_not_the_url():
     edited = json.loads(_real_map_bytes())

tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py

@@ -4114,6 +4114,35 @@ class TestModelInfoCostMapEchoFilter:
         info = json.loads(result["model_info"])
         assert "max_input_tokens" not in info, info
 
+    def test_reset_to_a_remote_catalog_value_that_differs_from_the_bundled_one(self, monkeypatch: pytest.MonkeyPatch):
+        from types import MappingProxyType
+
+        import litellm
+
+        from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
+        from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
+        from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
+
+        bundled = litellm.get_model_info("openai/gpt-5.6")
+        remote = {**bundled, "max_input_tokens": bundled["max_input_tokens"] + 1}
+        monkeypatch.setattr(
+            GetModelCostMap, "_loaded_catalog", MappingProxyType({remote["key"]: MappingProxyType(remote)})
+        )
+        monkeypatch.setattr(litellm, "get_model_info", lambda model, **_: {**remote, "max_input_tokens": 2048})
+        db_model = Deployment(
+            model_name="gpt-5.6",
+            litellm_params=LiteLLM_Params(model="openai/gpt-5.6"),
+            model_info=ModelInfo(id="dep-echo-9", max_input_tokens=2048),
+        )
+
+        result = update_db_model(
+            db_model=db_model,
+            updated_patch=updateDeployment(model_info=ModelInfo(**{**remote, "id": "dep-echo-9", "db_model": True})),
+        )
+
+        info = json.loads(result["model_info"])
+        assert "max_input_tokens" not in info, info
+
     def test_echo_is_compared_against_the_deployments_lookup_not_the_key(self):
         import litellm