From 3f712e3fde106b5da383f5a649a3fe04289f03e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 15:49:19 -0700 Subject: [PATCH] fix(router): stop custom model_info leaking onto shared backend cost map key --- litellm/router.py | 27 ++-- litellm/types/utils.py | 16 ++ .../test_router_model_cost_isolation.py | 138 ++++++++++++++++++ 3 files changed, 169 insertions(+), 12 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index ae3f7ba11c2..3ecaef591f3 100644 --- a/litellm/router.py +++ b/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/", "") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ec8a9336ca7..5b98e8be8d2 100644 --- a/litellm/types/utils.py +++ b/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 diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 6db7b04b3b7..c7f5513b94b 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/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