From ad4e6e2395620ea6d2fe38089a54cde160720de2 Mon Sep 17 00:00:00 2001 From: Aarkin Karnik <56022539+Aarkin7@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:20:02 +0530 Subject: [PATCH] fix(router): clean pattern_router state on upsert/delete (#29601) * fix(router): clean pattern_router state on upsert/delete PatternMatchRouter.add_pattern was append-only, and neither Router.upsert_deployment nor Router.delete_deployment removed the existing entry. Rotated-out api_keys stayed in the routing rotation for wildcard deployments (model_name with `*`) until proxy restart, silently defeating key rotation as an admin operation. The same leak applied to provider_default_deployment_ids and per-team pattern routers, and the patterns list grew unboundedly on every edit * test(router): direct unit tests for _remove_deployment_from_wildcard_state router_code_coverage.py greps test files for AST Call nodes and flagged the helper as untested because the existing coverage only exercised it transitively through upsert/delete. Adds two direct tests that pin the helper's contract (cleans across global pattern router, per-team routers with empty-router pop, and provider_default_deployment_ids; noop on falsy model_id) * fix(router): address Greptile review on pattern_router cleanup Widen PatternMatchRouter.remove_deployment annotation to Optional[str]; the implementation already handles None via the falsy guard and the unit test exercises it directly. Move _remove_deployment_from_wildcard_state up one level in upsert_deployment so it runs whenever the prior deployment is on the router, not only when the model_id is present in the fast-mapping index. The scenario is currently unreachable (get_deployment shares the same index), but the cleanup is idempotent so this is defensive against any future divergence between those code paths. * fix(router): widen _remove_deployment_from_wildcard_state to Optional[str] Moving the call out of the inner `deployment_id in deployment_fast_mapping` block in the previous commit lost mypy's narrowing of `deployment_id` from Optional[str] to str, tripping the lint CI. The helper already handles None via its falsy guard, so widening the annotation matches the actual contract. * fix(router): make delete_deployment wildcard cleanup symmetric with upsert After the previous commit moved _remove_deployment_from_wildcard_state out of the inner index-map guard in upsert_deployment, delete_deployment was still calling it only inside `if deployment_idx is not None`. Greptile flagged the asymmetry: under a desynced index_map, delete would silently leave the stale wildcard credential in pattern_router. Moves the cleanup call to the top of the try block, mirroring the upsert path. Cleanup is idempotent so the change is a no-op on the happy path. Adds a regression test that simulates the desync by removing the entry from model_id_to_deployment_index_map and asserts delete still clears pattern_router. --- litellm/router.py | 27 ++ .../router_utils/pattern_match_deployments.py | 28 +++ .../test_router_pattern_matching.py | 87 +++++++ tests/test_litellm/test_router.py | 238 ++++++++++++++++++ 4 files changed, 380 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index 80584858311..9fce6d3e897 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8855,6 +8855,8 @@ class Router: model_id=deployment_id, removal_idx=removal_idx ) + self._remove_deployment_from_wildcard_state(model_id=deployment_id) + # if the model_id is not in router self.add_deployment(deployment=deployment) return deployment @@ -8881,6 +8883,9 @@ class Router: deployment_idx = self.model_id_to_deployment_index_map[id] try: + # Idempotent and symmetric with upsert_deployment, so a desynced + # index_map cannot leave stale wildcard credentials behind. + self._remove_deployment_from_wildcard_state(model_id=id) if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) @@ -8898,6 +8903,28 @@ class Router: except Exception: return None + def _remove_deployment_from_wildcard_state(self, model_id: Optional[str]) -> None: + """ + Drop every reference to model_id from the wildcard-routing data + structures. Without this, upsert/delete leaves stale credentials in + pattern_router which silently defeats key rotation for wildcard + deployments. + """ + if not model_id: + return + self.pattern_router.remove_deployment(model_id) + empty_team_ids: List[str] = [] + for team_id, team_router in self.team_pattern_routers.items(): + team_router.remove_deployment(model_id) + if not team_router.patterns: + empty_team_ids.append(team_id) + for team_id in empty_team_ids: + del self.team_pattern_routers[team_id] + if model_id in self.provider_default_deployment_ids: + self.provider_default_deployment_ids = [ + i for i in self.provider_default_deployment_ids if i != model_id + ] + def _get_router_deployment_budget_limiter( self, ) -> Optional[RouterBudgetLimiting]: diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 48f85a83411..41bb0061784 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -75,6 +75,34 @@ class PatternMatchRouter: self.patterns[regex] = [] self.patterns[regex].append(llm_deployment) + def remove_deployment(self, model_id: Optional[str]) -> int: + """ + Remove every stored entry whose model_info.id equals model_id. + + Returns the number of entries removed. A regex whose deployment list + becomes empty is dropped so the patterns dict does not grow unboundedly. + + Empty / falsy model_id is a no-op so callers cannot accidentally wipe + entries whose model_info.id is missing. + """ + if not model_id: + return 0 + removed_total = 0 + for regex in list(self.patterns.keys()): + original = self.patterns[regex] + filtered = [ + d for d in original if (d.get("model_info") or {}).get("id") != model_id + ] + removed_here = len(original) - len(filtered) + if removed_here == 0: + continue + removed_total += removed_here + if filtered: + self.patterns[regex] = filtered + else: + del self.patterns[regex] + return removed_total + def _pattern_to_regex(self, pattern: str) -> str: """ Convert a wildcard pattern to a regex pattern diff --git a/tests/local_testing/test_router_pattern_matching.py b/tests/local_testing/test_router_pattern_matching.py index d09790d43b1..ddea3cb2a15 100644 --- a/tests/local_testing/test_router_pattern_matching.py +++ b/tests/local_testing/test_router_pattern_matching.py @@ -395,3 +395,90 @@ def test_wildcard_priority_over_deployment_names(): assert ( deployments[0]["litellm_params"]["api_base"] == "http://localhost:8081/openai" ), f"Expected '*' wildcard deployment (8081), got {deployments[0]['litellm_params']['api_base']}" + + +def _make_wildcard_entry(model_id: str, api_key: str) -> dict: + """Build a dict shaped like Deployment(...).to_json(exclude_none=True).""" + return Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params(model="openai/*", api_key=api_key), + model_info=ModelInfo(id=model_id), + ).to_json(exclude_none=True) + + +def test_remove_deployment_drops_matching_id_only(): + router = PatternMatchRouter() + router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A")) + router.add_pattern("openai/*", _make_wildcard_entry("id-B", "key-B")) + + removed = router.remove_deployment("id-A") + + assert removed == 1 + survivors = router.patterns["openai/(.*)"] + assert len(survivors) == 1 + assert survivors[0]["model_info"]["id"] == "id-B" + assert survivors[0]["litellm_params"]["api_key"] == "key-B" + + +def test_remove_deployment_drops_regex_key_when_list_empties(): + router = PatternMatchRouter() + router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A")) + + router.remove_deployment("id-A") + + assert router.patterns == {} + + +def test_remove_deployment_spans_multiple_regexes(): + router = PatternMatchRouter() + router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A")) + router.add_pattern("anthropic/*", _make_wildcard_entry("id-A", "key-A")) + router.add_pattern("openai/*", _make_wildcard_entry("id-B", "key-B")) + + removed = router.remove_deployment("id-A") + + assert removed == 2 + assert "anthropic/(.*)" not in router.patterns + assert [d["model_info"]["id"] for d in router.patterns["openai/(.*)"]] == ["id-B"] + + +def test_remove_deployment_noop_for_unknown_id(): + router = PatternMatchRouter() + router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A")) + before = {k: list(v) for k, v in router.patterns.items()} + + removed = router.remove_deployment("id-DOES-NOT-EXIST") + + assert removed == 0 + assert router.patterns == before + + +def test_remove_deployment_with_falsy_id_is_noop_even_when_entries_have_no_id(): + """ + An entry without model_info.id has '' / None as its id. remove_deployment('') + must NOT match those entries; otherwise a stray empty-string call would + wipe every id-less wildcard deployment in the router. + """ + router = PatternMatchRouter() + router.patterns["openai/(.*)"] = [ + {"model_name": "openai/*", "litellm_params": {"api_key": "key-X"}} + ] + + for falsy in ("", None): + removed = router.remove_deployment(falsy) + assert removed == 0 + assert len(router.patterns["openai/(.*)"]) == 1 + + +def test_remove_deployment_tolerates_missing_model_info(): + router = PatternMatchRouter() + router.patterns["openai/(.*)"] = [ + {"model_name": "openai/*", "litellm_params": {"api_key": "X"}}, + _make_wildcard_entry("id-A", "key-A"), + ] + + router.remove_deployment("id-A") + + survivors = router.patterns["openai/(.*)"] + assert len(survivors) == 1 + assert "model_info" not in survivors[0] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 830edf6412d..6a921983a74 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4756,3 +4756,241 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) is True ) + + +def _build_wildcard_router(api_key: str, model_id: str = "wildcard-deployment-1"): + from litellm.router import Deployment + + router = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": api_key}, + "model_info": {"id": model_id}, + }, + ], + ) + return router, Deployment + + +def test_upsert_wildcard_deployment_removes_stale_api_key(): + """ + Regression: rotating the api_key on a wildcard deployment must NOT leave + the previous credential in pattern_router. Before the fix, the old entry + survived in pattern_router.patterns[regex] and was still selected on + roughly half of wildcard traffic. + """ + router, Deployment = _build_wildcard_router(api_key="OLD-KEY") + model_id = "wildcard-deployment-1" + + router.upsert_deployment( + Deployment( + model_name="openai/*", + litellm_params={"model": "openai/*", "api_key": "NEW-KEY"}, + model_info={"id": model_id}, + ) + ) + + entries = router.pattern_router.patterns["openai/(.*)"] + assert ( + len(entries) == 1 + ), f"pattern_router accumulated stale entries: {len(entries)} present" + assert entries[0]["litellm_params"]["api_key"] == "NEW-KEY" + assert all( + e["litellm_params"]["api_key"] != "OLD-KEY" for e in entries + ), "Rotated-out api_key is still resident in pattern_router" + + +def test_upsert_wildcard_deployment_is_idempotent_in_pattern_router(): + """ + Repeated upserts of the same wildcard deployment must not grow the + pattern_router unboundedly. + """ + router, Deployment = _build_wildcard_router(api_key="key-0") + model_id = "wildcard-deployment-1" + + for i in range(1, 6): + router.upsert_deployment( + Deployment( + model_name="openai/*", + litellm_params={"model": "openai/*", "api_key": f"key-{i}"}, + model_info={"id": model_id}, + ) + ) + + entries = router.pattern_router.patterns["openai/(.*)"] + assert len(entries) == 1 + assert entries[0]["litellm_params"]["api_key"] == "key-5" + assert router.provider_default_deployment_ids.count(model_id) == 1 + + +def test_delete_wildcard_deployment_clears_pattern_router(): + router, _ = _build_wildcard_router(api_key="OLD-KEY") + model_id = "wildcard-deployment-1" + + assert router.pattern_router.patterns # sanity + + router.delete_deployment(id=model_id) + + assert router.pattern_router.patterns == {} + assert model_id not in router.provider_default_deployment_ids + + +def test_upsert_wildcard_to_concrete_model_removes_old_wildcard_entry(): + """ + Changing model_name from a wildcard to a concrete name must drop the + old wildcard entry. Otherwise traffic to openai/ still routes + to a model the operator believes they renamed. + """ + router, Deployment = _build_wildcard_router(api_key="key-A") + model_id = "wildcard-deployment-1" + + router.upsert_deployment( + Deployment( + model_name="openai/gpt-4o", + litellm_params={"model": "openai/gpt-4o", "api_key": "key-A"}, + model_info={"id": model_id}, + ) + ) + + assert router.pattern_router.patterns == {} + assert model_id not in router.provider_default_deployment_ids + + +def test_upsert_team_wildcard_deployment_does_not_leak_old_key(): + from litellm.router import Deployment + + team_id = "team-acme" + model_id = "team-wildcard-1" + + router = litellm.Router( + model_list=[ + { + "model_name": "openai/team-foo-*", + "litellm_params": {"model": "openai/*", "api_key": "OLD-KEY"}, + "model_info": { + "id": model_id, + "team_id": team_id, + "team_public_model_name": "team-foo-*", + }, + }, + ], + ) + + assert team_id in router.team_pattern_routers + assert router.team_pattern_routers[team_id].patterns + + router.upsert_deployment( + Deployment( + model_name="openai/team-foo-*", + litellm_params={"model": "openai/*", "api_key": "NEW-KEY"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": "team-foo-*", + }, + ) + ) + + team_router = router.team_pattern_routers[team_id] + assert len(team_router.patterns) == 1 + team_entries = next(iter(team_router.patterns.values())) + assert len(team_entries) == 1 + assert team_entries[0]["litellm_params"]["api_key"] == "NEW-KEY" + + +def test_delete_team_wildcard_removes_empty_team_router(): + """ + When the last wildcard deployment for a team is deleted, its dedicated + team_pattern_router entry should be discarded too, not left as an empty + PatternMatchRouter holding the team_id forever. + """ + team_id = "team-acme" + model_id = "team-wildcard-1" + + router = litellm.Router( + model_list=[ + { + "model_name": "openai/team-foo-*", + "litellm_params": {"model": "openai/*", "api_key": "key-A"}, + "model_info": { + "id": model_id, + "team_id": team_id, + "team_public_model_name": "team-foo-*", + }, + }, + ], + ) + + router.delete_deployment(id=model_id) + + assert team_id not in router.team_pattern_routers + + +def test_remove_deployment_from_wildcard_state_cleans_all_three_structures(): + """ + Pin the contract of Router._remove_deployment_from_wildcard_state: a single + call must strip model_id from the global pattern router, from every per-team + pattern router (dropping the team entry when its router empties out), and + from provider_default_deployment_ids. Exercised directly so the helper's + behavior is locked in independently of the upsert/delete code paths. + """ + from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + + router, _ = _build_wildcard_router(api_key="key-A") + model_id = "wildcard-deployment-1" + + router.team_pattern_routers["team-x"] = PatternMatchRouter() + router.team_pattern_routers["team-x"].add_pattern( + "openai/team-x-*", + { + "model_name": "openai/team-x-*", + "litellm_params": {"model": "openai/*", "api_key": "key-A"}, + "model_info": {"id": model_id}, + }, + ) + assert router.pattern_router.patterns + assert "team-x" in router.team_pattern_routers + assert model_id in router.provider_default_deployment_ids + + router._remove_deployment_from_wildcard_state(model_id=model_id) + + assert router.pattern_router.patterns == {} + assert "team-x" not in router.team_pattern_routers + assert model_id not in router.provider_default_deployment_ids + + +def test_remove_deployment_from_wildcard_state_is_noop_for_empty_id(): + """ + A falsy model_id must not touch any wildcard-routing state; otherwise an + accidental empty-string call could wipe deployments that lack model_info.id. + """ + router, _ = _build_wildcard_router(api_key="key-A") + snapshot_patterns = {k: list(v) for k, v in router.pattern_router.patterns.items()} + snapshot_ids = list(router.provider_default_deployment_ids) + + router._remove_deployment_from_wildcard_state(model_id="") + + assert router.pattern_router.patterns == snapshot_patterns + assert router.provider_default_deployment_ids == snapshot_ids + + +def test_delete_deployment_cleans_wildcard_state_even_when_index_is_desynced(): + """ + Symmetric with the upsert path: delete_deployment must clean wildcard + state from the model_id regardless of whether the fast-mapping index + still knows about it. Simulates index corruption / partial-failure by + removing the entry from model_id_to_deployment_index_map while leaving + pattern_router intact, then calls delete and asserts pattern_router is + cleaned. + """ + router, _ = _build_wildcard_router(api_key="key-A") + model_id = "wildcard-deployment-1" + + assert router.pattern_router.patterns + del router.model_id_to_deployment_index_map[model_id] + + router.delete_deployment(id=model_id) + + assert router.pattern_router.patterns == {} + assert model_id not in router.provider_default_deployment_ids