From eb142b900e1e2900b9f8ebf059c6ca0e95014f28 Mon Sep 17 00:00:00 2001 From: shivam Date: Tue, 12 May 2026 17:44:05 -0700 Subject: [PATCH] test(proxy): drop allow_client_tags opt-in gate and add credential rename cascade tests Removes the allow_client_tags metadata check from apply_client_tag_policy_pre_auth so x-litellm-tags headers are always merged into request metadata, matching the post-auth behavior in add_litellm_data_to_request. Updates pre-call tests accordingly and adds a new test suite covering cascading credential renames into model rows. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/litellm_pre_call_utils.py | 42 ++--- .../proxy/credential_endpoints/__init__.py | 0 .../test_credential_rename_cascade.py | 165 ++++++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 76 ++------ 4 files changed, 194 insertions(+), 89 deletions(-) create mode 100644 tests/test_litellm/proxy/credential_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/credential_endpoints/test_credential_rename_cascade.py diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 78e63f4094..9f6e21521c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1197,44 +1197,24 @@ class LiteLLMProxyRequestSetup: user_api_key_dict: UserAPIKeyAuth, ) -> None: """ - Apply the client-tag policy BEFORE auth budget gates run, so - ``_tag_max_budget_check`` (which only inspects ``request_data``) - sees ``x-litellm-tags`` header tags. Without this, header-tagged + Merge ``x-litellm-tags`` header tags into ``request_data`` BEFORE + auth budget gates run, so ``_tag_max_budget_check`` (which only + inspects ``request_data``) sees them. Without this, header-tagged requests silently bypass per-tag budget enforcement. - Mirrors the strip + merge that ``add_litellm_data_to_request`` - performs post-auth, gated on the same ``allow_client_tags`` flag. - - Why: ``add_litellm_data_to_request`` runs after the auth chain has - completed, so any header-supplied tags it merges in are invisible - to ``_tag_max_budget_check``. Running the merge here closes that - gap. The post-auth strip + merge remains as defense-in-depth. + Why: ``add_litellm_data_to_request`` runs the equivalent merge + post-auth, after ``_tag_max_budget_check`` has already executed. + Header-supplied tags merged there are invisible to that check. + Running the merge here closes that gap; the post-auth merge in + ``add_litellm_data_to_request`` remains as defense-in-depth. How to apply: invoked from the auth chain just before ``common_checks``. Mutates ``request_data`` in place; idempotent when followed by ``add_litellm_data_to_request``. """ - _admin_allow_client_tags = False - for _admin_meta in ( - user_api_key_dict.metadata, - user_api_key_dict.team_metadata, - ): - if ( - isinstance(_admin_meta, dict) - and _admin_meta.get("allow_client_tags") is True - ): - _admin_allow_client_tags = True - break - - if not _admin_allow_client_tags: - # Don't strip body-supplied tags here — pre-PR behavior was that - # _tag_max_budget_check (inside common_checks) saw and enforced - # per-tag budgets on body tags regardless of allow_client_tags. - # Stripping pre-auth would silently disable that enforcement. - # The post-auth strip in add_litellm_data_to_request still - # removes unauthorized tags before they leave the proxy. - return - + # No allow_client_tags opt-in: caller-supplied tags always flow + # into metadata.tags (see add_litellm_data_to_request). The pre-auth + # merge mirrors that so _tag_max_budget_check sees the same tags. headers = _safe_get_request_headers(request=request) raw_header_tags = headers.get("x-litellm-tags") if not raw_header_tags: diff --git a/tests/test_litellm/proxy/credential_endpoints/__init__.py b/tests/test_litellm/proxy/credential_endpoints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/proxy/credential_endpoints/test_credential_rename_cascade.py b/tests/test_litellm/proxy/credential_endpoints/test_credential_rename_cascade.py new file mode 100644 index 0000000000..e8d0e04876 --- /dev/null +++ b/tests/test_litellm/proxy/credential_endpoints/test_credential_rename_cascade.py @@ -0,0 +1,165 @@ +""" +Tests for cascading credential renames into model rows. + +When a credential is renamed via PATCH /credentials/{old_name}, every model +row whose `litellm_params.litellm_credential_name` references the old name +must be updated in lockstep — otherwise those models will fail at request +time when the router tries to resolve a credential that no longer exists. +""" + +import json +import types +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm.proxy.proxy_server as ps +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.credential_endpoints.endpoints import ( + _cascade_rename_credential_in_models, +) + + +@pytest.fixture +def salt_key(monkeypatch): + """Encrypt/decrypt helpers require a signing key — set one for the test.""" + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-key-cascade-rename") + yield + + +def _model_row(model_id: str, credential_name_plain: str | None): + """ + Build a fake LiteLLM_ProxyModelTable row whose `litellm_params` mirrors the + on-disk shape (encrypted values, dict-typed JSON column). + """ + params: dict = {"model": "gpt-4o"} + if credential_name_plain is not None: + params["litellm_credential_name"] = encrypt_value_helper( + value=credential_name_plain + ) + row = MagicMock() + row.model_id = model_id + row.litellm_params = params + return row + + +@pytest.mark.asyncio +async def test_cascade_renames_only_matching_models(salt_key, monkeypatch): + """Only models referencing the old name are updated; others are left alone.""" + monkeypatch.setattr(ps, "llm_router", None) + + matching = _model_row("model-1", "old-cred") + other = _model_row("model-2", "different-cred") + no_credential = _model_row("model-3", None) + + tx = types.SimpleNamespace( + litellm_proxymodeltable=types.SimpleNamespace( + find_many=AsyncMock(return_value=[matching, other, no_credential]), + update=AsyncMock(), + ) + ) + + updated = await _cascade_rename_credential_in_models( + tx=tx, + old_credential_name="old-cred", + new_credential_name="new-cred", + ) + + assert updated == 1 + tx.litellm_proxymodeltable.update.assert_awaited_once() + call = tx.litellm_proxymodeltable.update.await_args + assert call.kwargs["where"] == {"model_id": "model-1"} + + # Prisma's JSON column expects a serialized string, not a dict. + raw_params = call.kwargs["data"]["litellm_params"] + assert isinstance(raw_params, str) + written_params = json.loads(raw_params) + assert written_params["model"] == "gpt-4o" + # The new credential name is stored encrypted, not plain text. + assert written_params["litellm_credential_name"] != "new-cred" + # And it must round-trip back to the new name. + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + ) + + assert ( + decrypt_value_helper( + value=written_params["litellm_credential_name"], + key="litellm_credential_name", + return_original_value=True, + ) + == "new-cred" + ) + + +@pytest.mark.asyncio +async def test_cascade_updates_in_memory_router(salt_key, monkeypatch): + """ + The router's in-memory model_list holds decrypted credential names. The + cascade must rewrite them so live traffic doesn't try to resolve a name + that no longer exists in litellm.credential_list. + """ + fake_router = MagicMock() + fake_router.model_list = [ + {"litellm_params": {"litellm_credential_name": "old-cred", "model": "gpt-4o"}}, + { + "litellm_params": { + "litellm_credential_name": "other-cred", + "model": "claude", + } + }, + {"litellm_params": {"model": "no-creds"}}, + ] + monkeypatch.setattr(ps, "llm_router", fake_router) + + matching = _model_row("model-1", "old-cred") + tx = types.SimpleNamespace( + litellm_proxymodeltable=types.SimpleNamespace( + find_many=AsyncMock(return_value=[matching]), + update=AsyncMock(), + ) + ) + + await _cascade_rename_credential_in_models( + tx=tx, + old_credential_name="old-cred", + new_credential_name="new-cred", + ) + + assert ( + fake_router.model_list[0]["litellm_params"]["litellm_credential_name"] + == "new-cred" + ) + assert ( + fake_router.model_list[1]["litellm_params"]["litellm_credential_name"] + == "other-cred" + ) + assert "litellm_credential_name" not in fake_router.model_list[2]["litellm_params"] + + +@pytest.mark.asyncio +async def test_cascade_noop_when_no_models_match(salt_key, monkeypatch): + """No matching rows → no updates issued, no in-memory mutation.""" + fake_router = MagicMock() + untouched = { + "litellm_params": {"litellm_credential_name": "other-cred", "model": "gpt-4o"} + } + fake_router.model_list = [untouched] + monkeypatch.setattr(ps, "llm_router", fake_router) + + tx = types.SimpleNamespace( + litellm_proxymodeltable=types.SimpleNamespace( + find_many=AsyncMock(return_value=[_model_row("model-2", "other-cred")]), + update=AsyncMock(), + ) + ) + + updated = await _cascade_rename_credential_in_models( + tx=tx, + old_credential_name="old-cred", + new_credential_name="new-cred", + ) + + assert updated == 0 + tx.litellm_proxymodeltable.update.assert_not_awaited() + assert untouched["litellm_params"]["litellm_credential_name"] == "other-cred" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index edc5aa1bec..18c11dcc20 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -3873,14 +3873,14 @@ class TestApplyClientTagPolicyPreAuth: post-auth in ``add_litellm_data_to_request``. """ - def test_merges_header_tags_into_metadata_when_opted_in(self): + def test_merges_header_tags_into_metadata(self): request_mock = _build_request_mock_with_headers( {"x-litellm-tags": "tenant:acme,env:prod"} ) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", - metadata={"allow_client_tags": True}, + metadata={}, team_metadata={}, ) @@ -3902,7 +3902,7 @@ class TestApplyClientTagPolicyPreAuth: } user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", - metadata={"allow_client_tags": True}, + metadata={}, team_metadata={}, ) @@ -3915,21 +3915,17 @@ class TestApplyClientTagPolicyPreAuth: # Existing tags first, dedupe header tags assert data["metadata"]["tags"] == ["env:prod", "team:platform", "tenant:acme"] - def test_preserves_body_tags_when_not_opted_in(self): - # Pre-auth must NOT strip body-supplied tags for non-opted-in keys. - # _tag_max_budget_check (inside common_checks) enforces per-tag - # budgets on whatever tags it sees in request_data, and pre-PR - # behavior was that body tags hit that check regardless of - # allow_client_tags. The post-auth strip in add_litellm_data_to_request - # cleans them up before they leave the proxy — that's covered by a - # separate regression test. + def test_preserves_body_tags(self): + # Pre-auth must NOT touch body-supplied tags. _tag_max_budget_check + # (inside common_checks) enforces per-tag budgets on whatever tags + # it sees in request_data, including body tags. The helper only + # adds header tags to metadata.tags. request_mock = _build_request_mock_with_headers( {"x-litellm-tags": "tenant:acme"} ) data = { "model": "gpt-3.5-turbo", "tags": ["root-tag"], - "metadata": {"tags": ["meta-tag"]}, "litellm_metadata": {"tags": ["litellm-meta-tag"]}, } user_api_key_dict = UserAPIKeyAuth( @@ -3945,48 +3941,12 @@ class TestApplyClientTagPolicyPreAuth: ) assert data["tags"] == ["root-tag"] - assert data["metadata"]["tags"] == ["meta-tag"] - assert data["litellm_metadata"]["tags"] == ["litellm-meta-tag"] - - def test_does_not_merge_header_tags_when_not_opted_in(self): - # Even with the header set, no opt-in means the header is ignored - # and metadata.tags is not created from it. - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) - data = {"model": "gpt-3.5-turbo"} - user_api_key_dict = UserAPIKeyAuth( - api_key="hashed-key", - metadata={}, - team_metadata={}, - ) - - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( - request=request_mock, - request_data=data, - user_api_key_dict=user_api_key_dict, - ) - - assert "tags" not in data.get("metadata", {}) - - def test_team_metadata_opt_in_is_honored(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) - data = {"model": "gpt-3.5-turbo"} - user_api_key_dict = UserAPIKeyAuth( - api_key="hashed-key", - metadata={}, - team_metadata={"allow_client_tags": True}, - ) - - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( - request=request_mock, - request_data=data, - user_api_key_dict=user_api_key_dict, - ) - - assert data["metadata"]["tags"] == ["tenant:acme"] + # litellm_metadata is the active metadata key (it's present), so + # header tags merge into it and union with existing tags there. + assert data["litellm_metadata"]["tags"] == [ + "litellm-meta-tag", + "tenant:acme", + ] def test_uses_litellm_metadata_when_present(self): request_mock = _build_request_mock_with_headers( @@ -3998,7 +3958,7 @@ class TestApplyClientTagPolicyPreAuth: } user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", - metadata={"allow_client_tags": True}, + metadata={}, team_metadata={}, ) @@ -4014,12 +3974,12 @@ class TestApplyClientTagPolicyPreAuth: assert data["litellm_metadata"]["tags"] == ["tenant:acme"] assert "tags" not in data.get("metadata", {}) - def test_no_header_no_mutation_when_opted_in(self): + def test_no_header_no_mutation(self): request_mock = _build_request_mock_with_headers({}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", - metadata={"allow_client_tags": True}, + metadata={}, team_metadata={}, ) @@ -4045,7 +4005,7 @@ class TestApplyClientTagPolicyPreAuth: data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", - metadata={"allow_client_tags": True}, + metadata={}, team_metadata={}, )