fix(register_model): preserve built-in cache pricing when registering custom overrides under unmapped keys

litellm/proxy/proxy_server.py

@@ -2263,7 +2263,7 @@ async def _reconcile_budget_reservation_for_counter_update(
         )
     except Exception:
         verbose_proxy_logger.warning(
-            "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and continuing",
+            "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and falling back to direct increment",
             exc_info=True,
         )
         try:
@@ -2274,6 +2274,7 @@ async def _reconcile_budget_reservation_for_counter_update(
             verbose_proxy_logger.exception(
                 "Failed to invalidate reserved counters after reservation reconciliation failed"
             )
+        return set()
     return reserved_counter_keys
 
 

litellm/router.py

@@ -7738,6 +7738,39 @@ def _generate_model_id(self, model_group: str, litellm_params: dict):
 
         return hash_object.hexdigest()
 
+    @staticmethod
+    def _inherit_builtin_cache_pricing(
+        model_info: dict, backend_model: str, custom_llm_provider: Optional[str]
+    ) -> None:
+        """Fill missing cache pricing on a custom-priced deployment entry from
+        the backend model's built-in cost map entry, so a deployment that
+        only spells out ``input_cost_per_token``/``output_cost_per_token``
+        does not silently bill cache_read/cache_creation at 0.
+
+        User-specified cache fields always win; only ``None``/missing entries
+        are inherited. No-op when the backend model has no canonical entry.
+        """
+        cache_fields = (
+            "cache_creation_input_token_cost",
+            "cache_creation_input_token_cost_above_1hr",
+            "cache_creation_input_token_cost_above_200k_tokens",
+            "cache_read_input_token_cost",
+            "cache_read_input_token_cost_above_200k_tokens",
+        )
+        if all(model_info.get(f) is not None for f in cache_fields):
+            return
+        try:
+            backend_info = litellm.get_model_info(
+                model=backend_model, custom_llm_provider=custom_llm_provider
+            )
+        except Exception:
+            return
+        for field in cache_fields:
+            if model_info.get(field) is None:
+                backend_value = backend_info.get(field)
+                if backend_value is not None:
+                    model_info[field] = backend_value
+
     def _create_deployment(
         self,
         deployment_info: dict,
@@ -7766,6 +7799,13 @@ def _create_deployment(
                 if deployment.litellm_params.get(field) is not None:
                     _model_info[field] = deployment.litellm_params[field]
 
+            if _model_info.get("input_cost_per_token") is not None:
+                Router._inherit_builtin_cache_pricing(
+                    model_info=_model_info,
+                    backend_model=deployment.litellm_params.model,
+                    custom_llm_provider=deployment.litellm_params.custom_llm_provider,
+                )
+
             ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
             model_id = deployment.model_info.id
             if model_id is not None:
@@ -8505,6 +8545,13 @@ def add_deployment(self, deployment: Deployment) -> Optional[Deployment]:
             if field_value is not None:
                 _model_info_dict[field] = field_value
 
+        if _model_info_dict.get("input_cost_per_token") is not None:
+            Router._inherit_builtin_cache_pricing(
+                model_info=_model_info_dict,
+                backend_model=deployment.litellm_params.model,
+                custom_llm_provider=deployment.litellm_params.custom_llm_provider,
+            )
+
         # Register custom pricing in litellm.model_cost.
         # Mirrors _create_deployment() logic to ensure dynamically-added deployments
         # (e.g., loaded from DB) also have their custom pricing registered.

litellm/utils.py

@@ -2887,6 +2887,61 @@ def _convert_stringified_numbers(value):
     return value
 
 
+_BEDROCK_REGION_PREFIXES = (
+    "us.",
+    "eu.",
+    "apac.",
+    "jp.",
+    "au.",
+    "us-gov.",
+    "global.",
+    "ap-northeast-1.",
+)
+
+_CACHE_PRICING_FIELDS = (
+    "cache_creation_input_token_cost",
+    "cache_creation_input_token_cost_above_1hr",
+    "cache_creation_input_token_cost_above_200k_tokens",
+    "cache_read_input_token_cost",
+    "cache_read_input_token_cost_above_200k_tokens",
+)
+
+
+def _resolve_builtin_model_cost_entry(
+    key: str, provider: str
+) -> Optional[Dict[str, Any]]:
+    """Best-effort lookup of a built-in ``model_cost`` entry for a custom key
+    whose shape ``get_model_info`` cannot resolve (double provider prefixes
+    like ``bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region aliases).
+
+    Returns a copy of the matching entry so the caller can inherit its defaults
+    (most importantly cache pricing) without mutating the shared built-in.
+    Returns ``None`` when no safe match exists.
+    """
+    candidates: List[str] = []
+    segments = key.split("/")
+    idx = 0
+    while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet:
+        idx += 1
+        candidates.append("/".join(segments[idx:]))
+
+    base = candidates[-1] if candidates else key
+    for region_prefix in _BEDROCK_REGION_PREFIXES:
+        if base.startswith(region_prefix):
+            candidates.append(base[len(region_prefix) :])
+
+    if provider:
+        stripped = _strip_model_name(model=base, custom_llm_provider=provider)
+        if stripped != base:
+            candidates.append(stripped)
+
+    for candidate in candidates:
+        entry = litellm.model_cost.get(candidate)
+        if entry is not None and entry.get("litellm_provider") is not None:
+            return dict(entry)
+    return None
+
+
 def register_model(model_cost: Union[str, dict]):  # noqa: PLR0915
     """
     Register new / Override existing models (and their pricing) to specific providers.
@@ -2933,6 +2988,26 @@ def register_model(model_cost: Union[str, dict]):  # noqa: PLR0915
             except Exception:
                 existing_model = {}
                 model_cost_key = key
+                builtin_entry = _resolve_builtin_model_cost_entry(
+                    key=_key_str, provider=provider
+                )
+                if builtin_entry is not None:
+                    for field in _CACHE_PRICING_FIELDS:
+                        if (
+                            value.get(field) is None
+                            and builtin_entry.get(field) is not None
+                        ):
+                            existing_model[field] = builtin_entry[field]
+                elif (
+                    value.get("cache_creation_input_token_cost") is None
+                    and value.get("cache_read_input_token_cost") is None
+                ):
+                    verbose_logger.warning(
+                        f"register_model: model={key} not in built-in cost map and no "
+                        "prefix/region variant matched; cache cost fields will default "
+                        "to 0. To track cache cost, add cache_creation_input_token_cost "
+                        "and cache_read_input_token_cost to model_info"
+                    )
         # ``get_model_info`` returns ``litellm_provider: None`` when the
         # provider is unknown (e.g. custom deployments registered via
         # ``Router.add_deployment``). Persisting that None into

tests/test_litellm/proxy/proxy_server/test_spend_counters.py

@@ -192,8 +192,9 @@ async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set
 async def test_reconcile_budget_reservation_for_counter_update_failure_invalidates(
     monkeypatch,
 ):
-    """Reservation reconcile raising must invalidate reserved counters but
-    not propagate the exception."""
+    """Reservation reconcile raising must invalidate reserved counters, swallow
+    the exception, and return an empty set so the caller falls back to the
+    direct spend-counter increment instead of skipping it."""
     import litellm.proxy.spend_tracking.budget_reservation as br
 
     monkeypatch.setattr(
@@ -213,7 +214,7 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat
         budget_reservation={"foo": "bar"}, response_cost=1.0
     )
 
-    assert result == {"spend:key:abc"}
+    assert result == set()
     assert fake_invalidate.called is True
 
 

tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py

@@ -0,0 +1,87 @@
+"""
+Regression test for enforced-spend underreporting when Redis fails during the
+budget-reservation reconcile step of ``increment_spend_counters``.
+
+Production failure mode: a managed Redis returns an intermittent timeout on the
+reconcile increment. Reconcile deletes (invalidates) the shared counter and
+gives up, but ``increment_spend_counters`` still treats the counter as
+"already reconciled" and skips the direct increment. The actual call cost never
+lands in the enforced counter, so budgets stop gating until the next cold
+reseed pulls a lagging value from the DB.
+
+The fix makes the reconcile path fall back to the direct increment when it
+fails, so the actual cost is always written to the shared counter.
+"""
+
+import pytest
+
+from litellm.caching import DualCache
+from litellm.proxy import proxy_server
+
+
+class _FlakyRedisCache:
+    def __init__(self) -> None:
+        self._store: dict = {}
+        self._increment_calls = 0
+
+    async def async_increment(self, key, value, **kwargs):
+        self._increment_calls += 1
+        if self._increment_calls == 1:
+            raise Exception("Redis timeout")
+        self._store[key] = float(self._store.get(key, 0.0)) + float(value)
+        return self._store[key]
+
+    async def async_get_cache(self, key, *args, **kwargs):
+        return self._store.get(key)
+
+    async def async_delete_cache(self, key, *args, **kwargs):
+        self._store.pop(key, None)
+
+    async def async_set_cache(self, key, value, *args, **kwargs):
+        self._store[key] = float(value)
+        return True
+
+
+@pytest.mark.asyncio
+async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure(
+    monkeypatch,
+):
+    hashed_token = "hashed_test_token"
+    counter_key = f"spend:key:{hashed_token}"
+    reserved_cost = 0.5
+    response_cost = 1.0
+
+    flaky_redis = _FlakyRedisCache()
+    flaky_redis._store[counter_key] = reserved_cost
+
+    monkeypatch.setattr(proxy_server, "prisma_client", None)
+    monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
+    monkeypatch.setattr(proxy_server.spend_counter_cache, "redis_cache", flaky_redis)
+    proxy_server.spend_counter_cache.in_memory_cache.set_cache(
+        key=counter_key, value=reserved_cost
+    )
+
+    budget_reservation = {
+        "reserved_cost": reserved_cost,
+        "finalized": False,
+        "entries": [
+            {
+                "counter_key": counter_key,
+                "entity_type": "Key",
+                "entity_id": hashed_token,
+                "reserved_cost": reserved_cost,
+                "applied_adjustment": 0.0,
+            }
+        ],
+    }
+
+    await proxy_server.increment_spend_counters(
+        token=hashed_token,
+        team_id=None,
+        user_id=None,
+        response_cost=response_cost,
+        budget_reservation=budget_reservation,
+    )
+
+    enforced_spend = await flaky_redis.async_get_cache(key=counter_key)
+    assert enforced_spend == response_cost

tests/test_litellm/proxy/test_proxy_server.py

@@ -6474,7 +6474,12 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation():
 
 
 @pytest.mark.asyncio
-async def test_increment_spend_counters_invalidates_bad_reserved_counter_without_failing():
+async def test_increment_spend_counters_falls_back_to_direct_increment_on_bad_reserved_counter():
+    """When the reservation reconcile fails, the reserved counters are
+    invalidated and the actual response cost must still be written via the
+    direct increment fallback. Leaving the counter at ``None`` lets the next
+    request reseed a stale value from the DB and silently stops budget gating,
+    which is the bug this fix addresses."""
     from litellm.caching.dual_cache import DualCache
     from litellm.proxy.proxy_server import increment_spend_counters
 
@@ -6515,7 +6520,7 @@ async def test_increment_spend_counters_invalidates_bad_reserved_counter_without
             counter_cache.in_memory_cache.get_cache(
                 key="spend:key:key-bad-reserved-counter"
             )
-            is None
+            == 0.25
         )
     finally:
         ps.spend_counter_cache = orig_counter