Merge pull request #34041 from BerriAI/litellm_lit4544_shared_key_model_info_leak

fix(router): stop custom model_info leaking onto shared backend cost map key

litellm/router.py

@@ -200,6 +200,7 @@ from litellm.types.utils import (
     CustomPricingLiteLLMParams,
     GenericBudgetConfigType,
     LiteLLMBatch,
+    shared_backend_model_info,
 )
 from litellm.types.utils import ModelInfo
 from litellm.types.utils import ModelInfo as ModelMapInfo
@@ -7495,12 +7496,13 @@ class Router:
             if deployment.litellm_params.custom_llm_provider is not None:
                 _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name
 
-            # For the shared backend key, strip custom pricing fields so that
-            # one deployment's pricing overrides don't pollute another
-            # deployment sharing the same backend model name.
-            # Each deployment's full pricing is already stored under its
-            # unique model_id above.
-            _shared_model_info = CustomPricingLiteLLMParams.strip_custom_pricing_fields(_model_info)
+            # For the shared backend key, keep only cost-map schema fields
+            # (minus custom pricing) so that one deployment's pricing overrides
+            # or custom metadata (id, access_via_team_ids, arbitrary keys)
+            # don't pollute another deployment sharing the same backend model
+            # name. Each deployment's full model_info is already stored under
+            # its unique model_id above.
+            _shared_model_info = shared_backend_model_info(_model_info)
             _existing_shared_mode = (cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {}).get("mode")
             _deployment_mode = _shared_model_info.get("mode")
             # Keep the built-in bridge mode stable for shared backend keys.
@@ -8219,12 +8221,13 @@ class Router:
         if deployment.litellm_params.custom_llm_provider is not None:
             _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name
 
-        # For the shared backend key, strip custom pricing fields so that
-        # one deployment's pricing overrides don't pollute another
-        # deployment sharing the same backend model name.
-        # Each deployment's full pricing is already stored under its
-        # unique model_id above (when present).
-        _shared_model_info = CustomPricingLiteLLMParams.strip_custom_pricing_fields(_model_info_dict)
+        # For the shared backend key, keep only cost-map schema fields
+        # (minus custom pricing) so that one deployment's pricing overrides
+        # or custom metadata (id, access_via_team_ids, arbitrary keys)
+        # don't pollute another deployment sharing the same backend model
+        # name. Each deployment's full model_info is already stored under
+        # its unique model_id above (when present).
+        _shared_model_info = shared_backend_model_info(_model_info_dict)
         _backend_alias_cost = {_model_name: _shared_model_info}
         if "responses/" in _model_name:
             _stripped_model_name = _model_name.replace("responses/", "")

litellm/types/utils.py

@@ -5,6 +5,7 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
+    FrozenSet,
     List,
     Literal,
     Mapping,
@@ -3112,6 +3113,21 @@ class CustomPricingLiteLLMParams(BaseModel):
         return {k: v for k, v in model_info.items() if k not in cls.model_fields}
 
 
+SHARED_BACKEND_MODEL_INFO_FIELDS: FrozenSet[str] = frozenset(
+    ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__
+) - frozenset(CustomPricingLiteLLMParams.model_fields)
+
+
+def shared_backend_model_info(model_info: Dict[str, Any]) -> Dict[str, Any]:
+    """Return only the fields safe to register under a shared ``{provider}/{model}``
+    key in ``litellm.model_cost``: cost-map schema fields (``ModelInfoBase``) minus
+    per-deployment pricing overrides. Per-deployment metadata (``id``,
+    ``access_via_team_ids``, arbitrary custom keys) never belongs on the shared key;
+    it stays under the deployment's unique model id.
+    """
+    return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS}
+
+
 # Server-controlled fields that bound or drive an interceptor's agentic loop
 # (depth, cycle fingerprints, ceiling, code-interpreter sandbox state). Listed
 # in all_litellm_params so they are treated as LiteLLM-level and excluded from

tests/test_litellm/test_router_model_cost_isolation.py

@@ -683,6 +683,144 @@ def test_custom_pricing_isolated_from_sibling_via_proxy_model_info_path():
         _restore_model_cost_entries(model_keys)
 
 
+def test_custom_model_info_metadata_not_leaked_to_shared_backend_key():
+    """LIT-4544: two deployments share the same backend model but carry
+    different custom model_info (arbitrary keys, access_via_team_ids, ids).
+    None of that per-deployment metadata may land on the shared backend key in
+    litellm.model_cost (served raw by /public/litellm_model_cost_map);
+    before the fix it was merged last-write-wins so values flipped randomly.
+    """
+    backend_model = "openai/gpt-4o-mini"
+    shared_keys = ("gpt-4o-mini", backend_model)
+    leak_fields = ("id", "additionalProp1", "access_via_team_ids", "db_model")
+
+    model_keys = {
+        key: copy.deepcopy(litellm.model_cost.get(key))
+        for key in (*shared_keys, "lit4544-deploy-a", "lit4544-deploy-b")
+    }
+    try:
+        Router(
+            model_list=[
+                {
+                    "model_name": "alias-unrestricted",
+                    "litellm_params": {
+                        "model": backend_model,
+                        "api_key": "fake-key-a",
+                    },
+                    "model_info": {
+                        "id": "lit4544-deploy-a",
+                        "additionalProp1": {"restricted": False, "model_location": "EU"},
+                    },
+                },
+                {
+                    "model_name": "alias-restricted",
+                    "litellm_params": {
+                        "model": backend_model,
+                        "api_key": "fake-key-b",
+                    },
+                    "model_info": {
+                        "id": "lit4544-deploy-b",
+                        "additionalProp1": {"restricted": True, "model_location": "US"},
+                        "access_via_team_ids": ["team-b-only"],
+                    },
+                },
+            ],
+        )
+
+        for shared_key in shared_keys:
+            shared_entry = litellm.model_cost.get(shared_key) or {}
+            leaked = [field for field in leak_fields if field in shared_entry]
+            assert not leaked, (
+                f"per-deployment metadata {leaked} leaked onto shared key "
+                f"{shared_key}: {shared_entry}"
+            )
+
+        entry_a = litellm.model_cost["lit4544-deploy-a"]
+        assert entry_a["additionalProp1"] == {"restricted": False, "model_location": "EU"}
+        entry_b = litellm.model_cost["lit4544-deploy-b"]
+        assert entry_b["additionalProp1"] == {"restricted": True, "model_location": "US"}
+        assert entry_b["access_via_team_ids"] == ["team-b-only"]
+    finally:
+        _restore_model_cost_entries(model_keys)
+
+
+def test_add_deployment_does_not_leak_custom_metadata_to_shared_backend_key():
+    """LIT-4544 dynamic path: deployments added at runtime (e.g. loaded from
+    the DB every scheduler cycle) must not re-pollute the shared backend key
+    with per-deployment metadata either.
+    """
+    backend_model = "openai/gpt-4o-mini"
+    shared_keys = ("gpt-4o-mini", backend_model)
+    deploy_id = "lit4544-add-deployment"
+
+    model_keys = {
+        key: copy.deepcopy(litellm.model_cost.get(key))
+        for key in (*shared_keys, deploy_id)
+    }
+    try:
+        router = Router(model_list=[])
+        router.add_deployment(
+            deployment=Deployment(
+                model_name="alias-dynamic",
+                litellm_params=LiteLLM_Params(
+                    model=backend_model,
+                    api_key="fake-key-dynamic",
+                ),
+                model_info=ModelInfo(
+                    id=deploy_id,
+                    additionalProp1={"restricted": True},
+                    access_via_team_ids=["team-dynamic"],
+                ),
+            )
+        )
+
+        for shared_key in shared_keys:
+            shared_entry = litellm.model_cost.get(shared_key) or {}
+            leaked = [
+                field
+                for field in ("id", "additionalProp1", "access_via_team_ids", "db_model")
+                if field in shared_entry
+            ]
+            assert not leaked, (
+                f"per-deployment metadata {leaked} leaked onto shared key "
+                f"{shared_key}: {shared_entry}"
+            )
+
+        assert litellm.model_cost[deploy_id]["access_via_team_ids"] == ["team-dynamic"]
+    finally:
+        _restore_model_cost_entries(model_keys)
+
+
+def test_shared_backend_model_info_keeps_schema_fields_and_drops_the_rest():
+    """Unit test of the whitelist helper: cost-map schema fields survive,
+    custom pricing overrides and per-deployment metadata do not.
+    """
+    from litellm.types.utils import shared_backend_model_info
+
+    filtered = shared_backend_model_info(
+        {
+            "mode": "chat",
+            "litellm_provider": "openai",
+            "max_tokens": 128000,
+            "supports_vision": True,
+            "input_cost_per_token": 0.99,
+            "output_cost_per_token": 0.99,
+            "id": "deploy-a",
+            "db_model": False,
+            "access_via_team_ids": ["team-a"],
+            "additionalProp1": {"restricted": True},
+            "base_model": "gpt-4o-mini",
+        }
+    )
+
+    assert filtered == {
+        "mode": "chat",
+        "litellm_provider": "openai",
+        "max_tokens": 128000,
+        "supports_vision": True,
+    }
+
+
 def test_wildcard_zero_cost_request_does_not_poison_named_deployment_pricing():
     """LIT-3991 end to end: a proxy has a named text-embedding-3-small
     deployment relying on built-in pricing plus an ``openai/*`` wildcard with