tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-01 16:21:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(proxy): drop allow_client_tags opt-in gate and add credential rename cascade tests

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
shivam
2026-05-12 17:44:05 -07:00
co-authored by Claude Opus 4.7
parent 9cad30a8df
commit eb142b900e
4 changed files with 194 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/litellm_pre_call_utils.py
+11 -31
View File
@@ -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:
tests/test_litellm/proxy/credential_endpoints/__init__.py
View File
tests/test_litellm/proxy/credential_endpoints/test_credential_rename_cascade.py
+165
View File
@@ -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"
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
+18 -58
View File
@@ -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={},
)