scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-17 04:34:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 172
litellm/tests/unit/proxy/auth/test_auth_checks.py
Mateo Wang 622daf3cdf
fix(scim): write a group's roster in one locked transaction instead of one member at a time (#45197) 
* fix(scim): write a group's roster in one locked transaction instead of one member at a time

POST, PUT and PATCH /scim/v2/Groups resolve every member in one query and write the roster, the
membership rows and the users' team tags in one transaction under the team's advisory lock, so a
500-member push runs a fixed number of statements instead of thousands. PUT now answers with the
roster it just wrote

* test(scim): assert the roster the IdP gets back and boot the bulk roster cell with two workers

* fix(scim): skip the team lock for a rename-only PATCH and lock member rows in id order before the roster write

A PATCH naming no members opened the roster transaction and took the team advisory lock for nothing. A SIGKILLed worker then left its orphaned query engine holding that lock until the 60 s transaction timeout, and /team/delete, which takes the same lock, timed out behind it. An empty delta is now answered from a plain team read, with no transaction and no lock

Two groups trading members in parallel locked user rows in opposite orders, and one request died with a Postgres deadlock. Every user row a push touches is now locked with SELECT ... ORDER BY user_id FOR UPDATE before the first write, so the two requests serialize

The bulk roster cell gets a 120 s client through a new owned_proxy_process client_timeout, so an unfixed push fails on the statement count instead of the harness clock, requires positive counts on both pushes, and gains a parallel roster-trade test that answered 500 before the lock

* fix(scim): evict the cached team on a group rename and read a no-op roster from the writer

A rename-only PUT or PATCH wrote the new alias and left the team cached under
its id and old alias on every worker, so keys kept logging spend against the
old team alias until the TTL expired. The group write now evicts the cached
team (and broadcasts the eviction) right after the alias or metadata update.

sync_team_roster reads the team from the writer when the plan names nobody,
so a replica cannot hand back the pre-rename row; a plan already matching the
roster returns inside the lock without writing, evicting, or emitting the
members metric; the response is built from the written row, so lastModified
reflects the roster write; and detaching a removed member bumps the user's
updated_at like every other user write.

* fix(scim): bump detached members' updated_at in UTC

The detach UPDATE wrote updated_at = CURRENT_TIMESTAMP, which Postgres casts
into Prisma's naive timestamp column in the session timezone, so on a database
whose session is not UTC the removed member's row landed hours away from the
UTC instants Prisma itself writes. Write NOW() AT TIME ZONE 'UTC' instead, and
have the bulk roster cell assert the bumped row sits between two reads of the
database's own UTC clock taken around the replace.

* fix(scim): chunk the group write's member lookups under the bind cap

Every Prisma "in" filter the group write issues bound one parameter per
member with nothing bounding the list, which check_unbounded_in_lists
reports as 11 new findings. The single-field lookups now go through
litellm.repositories.chunked_in (caller_transaction inside the roster
transaction, per_chunk_ok for the role recompute), and the one
cross-field OR read slices its values at IN_LIST_CHUNK_SIZE and records
the bound. A push of IN_LIST_CHUNK_SIZE + 1 members now reads in two
slices, and a member named only in the last slice still resolves.

* fix(proxy): send bulk cache evictions in slices so the broadcaster drops none

Veria pointed out that the invalidation publisher caps its pending tasks at
1,024 and drops the excess with a warning, and that a group push now detaches
every removed member's keys in one call instead of one member at a time, so a
removal of more than roughly 340 members (three cache keys each) silently left
stale key objects on the other replicas. The merge base never hit the cap
because each member was a separate request.

Bulk evictions (team, user, organization, and session key deletes, plus the
roster settle) now publish in slices of 256 and wait, bounded at two seconds,
for the backlog already handed to the event loop before each slice. A Redis
that never answers costs a bulk caller at most that wait per slice instead of
wedging the request, and single-key callers are untouched. Regression tests
push 1,500 keys and a 500-member removal through a slow publisher and assert
every eviction reaches Redis; both fail against the previous tip.

* fix(proxy): clear this worker's cache before pacing bulk eviction broadcasts

Greptile pointed out that pacing the broadcasts in slices also paced the
local eviction: a bulk delete waited for the publish backlog before dropping
the next slice from this worker's own memory, so on a Redis that accepts
connections but never answers a key deleted from the database kept
authenticating on the handling worker for up to two seconds per slice. The
merge base cleared every key locally at once.

Bulk evictions now drop every affected entry from this worker's memory first,
through the cache's own partition lookup since key objects live apart from
the other auth objects, and only then publish in paced slices. The roster
settle does the same for the deleted tokens, the JWT mapping keys and every
changed member's user and membership entries before any step that can wait.
Regression tests run both paths against a publisher that never answers and
check that nothing is still cached on this worker before the first bound
elapses; both fail with the local eviction removed.

* fix(proxy): drop the Redis copies of bulk-evicted auth objects before pacing the broadcasts

* docs(proxy): drop the module docstring that repeated the PR description

* fix(scim): bound bulk cache evictions to one Redis round trip and close lock-order gaps

Bulk key evictions now drop entries from memory, delete them from Redis in one
chunked DEL, and broadcast every key in one pipelined publish, so a large roster
removal no longer overflows the publisher's pending cap or costs a Redis round
trip per member. A push waits for in-flight membership loads before evicting.

User rows are locked in byte order to match the spend writer, a bulk user delete
takes every team lock before its first user write, and a POST whose roster write
fails deletes the empty team so the IdP retry creates it again.

* test(team): observe delete broadcasts at the Redis layer

* fix(scim): lock user rows in id order on every multi-row writer and create unprovisioned members concurrently

* fix(team): chunk the email lookup behind the member row lock

* fix(scim): create a member deleted mid-push again instead of failing the group write

* fix(scim): refuse strict-mode mid-push deletes with 400 and check license seats once per batch

A member deleted mid-push is created again only when scim_upsert_user allows it,
under the verb's own created_via. New group members are checked against the
license seat limit as one batch before any is created. Auth cache publish drops
log the key count instead of the key.

* fix(scim): lock member rows while recomputing admin roles after a group push

* fix(users): re-read teams under the row locks in bulk user delete so a racing group push keeps no deleted member

* fix(users): take /user/delete and SCIM DELETE /Users users off their current teams under the row locks so a racing group push keeps no deleted member

* fix(scim): keep retrying a push while deletes race it, lock rows in member bulk delete, evict removed keys from Redis first

A group push gave up after one retry when SCIM user deletes kept landing between attempts, answering 500 in 9 of 10 race rounds. It now retries while each attempt finds newly deleted members, and gives up only when a member it created again goes missing a second time.

The team member bulk delete now locks the removed members' user rows in the shared order and takes the team off their teams arrays in one array_remove statement, so it no longer writes back a teams array read before a concurrent push added another team.

A push that removes members now drops their deleted team keys from memory and Redis together with the members' other cache entries before the team eviction, so a removed key cannot refill from Redis mid-cleanup.

* fix(user-delete): keep the team alias label on the team members gauge after removals

* test(integration): count relay statements over plaintext so an SSL-enabled Postgres does not hide them

* fix(user-delete): skip a team deleted while the user delete waited on its lock instead of answering 404

* fix(scim): answer 404 and roll back when the team is deleted during a roster write

* test(integration): audit cells for SCIM bulk roster writes, user delete locking and the eviction broadcast
2026-10-15 10:57:00 -07:00

2382 lines
80 KiB
Python
Raw Blame History

# What is this?
## Tests if 'get_end_user_object' works as expected
import sys, os, asyncio, time, random, uuid
import hashlib
import json
import traceback
from dotenv import load_dotenv
load_dotenv()
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from copy import deepcopy
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest, litellm
import httpx
from fastapi import Request, status
from prisma import Prisma
from litellm._service_logger import ServiceTypes
from litellm.proxy._types import LiteLLM_OrganizationTable, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import get_org_object, get_user_object
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.caching.caching import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy._types import (
LiteLLM_EndUserTable,
LiteLLM_BudgetTable,
LiteLLM_UserTable,
LiteLLM_TeamTable,
Litellm_EntityType,
LitellmUserRoles,
ProxyException,
)
from litellm.proxy.utils import PrismaClient
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.auth.auth_checks import (
can_team_access_model,
delete_cache_key_objects,
is_model_cost_zero,
virtual_key_soft_budget_check,
_team_soft_budget_check,
common_checks,
)
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.utils import CallInfo
from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER
async def _run_team_trace_id_check(
*,
route: str,
method: str = "POST",
request_body: dict[str, object] | None = None,
headers: Mapping[str, str] | None = None,
team_metadata: Mapping[str, object] | None = None,
team_object: LiteLLM_TeamTable | None = None,
registered_pass_through_endpoint: bool = False,
) -> bool:
request_payload: Final[dict[str, object]] = {} if request_body is None else request_body
raw_headers: Final = tuple(
(name.lower().encode("latin-1"), value.encode("latin-1"))
for name, value in (() if headers is None else headers.items())
)
request: Final = Request(
{
"type": "http",
"method": method,
"path": route,
"headers": raw_headers,
"query_string": b"",
}
)
if registered_pass_through_endpoint:
endpoint: Final = SimpleNamespace()
setattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)
request.scope["endpoint"] = endpoint
team: Final = (
team_object
if team_object is not None
else LiteLLM_TeamTable(
team_id="trace-id-team",
metadata=None if team_metadata is None else dict(team_metadata),
)
)
valid_token: Final = UserAPIKeyAuth(token="test-token", team_id=team.team_id)
user: Final = LiteLLM_UserTable(user_id="trace-id-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
return await common_checks(
request_body=request_payload,
team_object=team,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(spec=ProxyLogging),
valid_token=valid_token,
request=request,
skip_budget_checks=True,
)
@pytest.mark.parametrize(
"route",
["/chat/completions", "/v1/messages", "/mcp/", "/a2a/my-agent"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_llm_mcp_and_agent_requests(route: str) -> None:
request_body: Final[dict[str, object]] = {"metadata": {"source": "unit-test"}}
original_body: Final = deepcopy(request_body)
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route=route,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
assert exc_info.value.message == (
"Team 'trace-id-team' requires a trace ID on every LLM, MCP and agent request. "
"Send an x-litellm-trace-id header. "
"LLM requests can also send a W3C traceparent header or set metadata.trace_id in the body."
)
assert request_body == original_body
@pytest.mark.parametrize(
("headers", "request_body"),
[
({"x-litellm-trace-id": "trace-123"}, {}),
({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}, {}),
({}, {"metadata": {"trace_id": "trace-123"}}),
({}, {"litellm_metadata": {"trace_id": "trace-123"}}),
({}, {"litellm_trace_id": "trace-123"}),
({}, {"litellm_trace_id": "", "metadata": {"trace_id": "trace-123"}}),
({}, {"litellm_trace_id": None, "metadata": {"trace_id": "trace-123"}}),
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_supported_trace_carriers(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
original_body: Final = deepcopy(request_body)
await _run_team_trace_id_check(
route="/chat/completions",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert request_body == original_body
@pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"])
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None:
assert await _run_team_trace_id_check(
route=route,
method="POST",
request_body={},
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat() -> None:
request_body: Final[dict[str, object]] = {
"metadata": {"trace_id": ""},
"litellm_metadata": {"trace_id": "trace-123"},
}
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/chat/completions",
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
"request_body",
[
{"litellm_metadata": {"trace_id": "trace-123"}},
{"metadata": {"trace_id": "trace-123"}},
],
ids=["selected-litellm-metadata", "promoted-requester-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route(
request_body: dict[str, object],
) -> None:
assert await _run_team_trace_id_check(
route="/v1/responses",
request_body=request_body,
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.parametrize(
("headers", "request_body"),
[
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{"litellm_trace_id": ""},
),
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{"litellm_trace_id": None},
),
],
ids=[
"traceparent-with-empty-top-level",
"traceparent-with-null-top-level",
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trace_id_is_explicit(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/chat/completions",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None:
assert await _run_team_trace_id_check(
route="/mcp/",
request_body={},
headers={"x-litellm-trace-id": "trace-mcp-123"},
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.parametrize(
("headers", "request_body"),
[
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{},
),
({}, {"metadata": {"trace_id": "trace-body-123"}}),
],
ids=["traceparent", "body-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_routes(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/anthropic/v1/messages",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None:
assert await _run_team_trace_id_check(
route="/anthropic/v1/messages",
request_body={},
headers={"x-litellm-trace-id": "trace-provider-123"},
team_metadata={"require_trace_id": True},
) is True
@pytest.mark.parametrize(
("headers", "request_body"),
[
(
{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
{},
),
({}, {"metadata": {"trace_id": "trace-body-123"}}),
],
ids=["traceparent", "body-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pass_through_routes(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/configured-passthrough/infer",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
registered_pass_through_endpoint=True,
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
("route", "request_body"),
[
("/mcp/", {}),
(
"/mcp-rest/tools/call",
{"server_id": "math", "name": "add", "arguments": {"left": 2, "right": 3}},
),
(
"/a2a/my-agent",
{
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
},
),
],
ids=["streamable-mcp", "mcp-rest", "a2a"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes(
route: str,
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route=route,
headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
("route", "request_body"),
[
(
"/mcp-rest/tools/call",
{
"server_id": "math",
"name": "add",
"arguments": {"left": 2, "right": 3},
"metadata": {"trace_id": "trace-body-123"},
},
),
(
"/a2a/my-agent",
{
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
"metadata": {"trace_id": "trace-body-123"},
},
),
],
ids=["mcp-rest", "a2a"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_body_metadata_on_mcp_and_agent_routes(
route: str,
request_body: dict[str, object],
) -> None:
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route=route,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
@pytest.mark.parametrize(
("headers", "request_body"),
[
({"traceparent": "invalid"}, {}),
({}, {"metadata": {"trace_id": ""}}),
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_invalid_trace_carriers(
headers: Mapping[str, str],
request_body: dict[str, object],
) -> None:
original_body: Final = deepcopy(request_body)
exc_info: Final[pytest.ExceptionInfo[ProxyException]]
with pytest.raises(ProxyException) as exc_info:
await _run_team_trace_id_check(
route="/chat/completions",
headers=headers,
request_body=request_body,
team_metadata={"require_trace_id": True},
)
assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
assert exc_info.value.param == "trace_id"
assert request_body == original_body
@pytest.mark.parametrize(
("route", "method", "team_metadata"),
[
("/v1/models", "GET", {"require_trace_id": True}),
("/chat/completions", "POST", None),
("/chat/completions", "POST", {"require_trace_id": False}),
("/key/generate", "POST", {"require_trace_id": True}),
],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_exempts_methods_routes_and_disabled_teams(
route: str,
method: str,
team_metadata: Mapping[str, object] | None,
) -> None:
assert await _run_team_trace_id_check(route=route, method=method, team_metadata=team_metadata) is True
@pytest.mark.asyncio
async def test_team_require_trace_id_skips_non_dict_team_metadata() -> None:
team: Final[LiteLLM_TeamTable] = LiteLLM_TeamTable.model_construct(
team_id="trace-id-team",
metadata="not a mapping",
)
assert await _run_team_trace_id_check(route="/v1/chat/completions", team_object=team) is True
@pytest.mark.parametrize("route", ("/lens/datasets", "/lens/tracing/keys"))
def test_lens_management_does_not_require_an_inference_user_parameter(route: str) -> None:
from starlette.requests import Request
from litellm.proxy.auth.auth_checks import _enforce_user_param_check
request: Final = Request({"type": "http", "method": "POST", "path": route})
_enforce_user_param_check({"enforce_user_param": True}, request, {}, route)
with pytest.raises(Exception, match="'user' param not passed"):
_enforce_user_param_check({"enforce_user_param": True}, request, {}, "/v1/chat/completions")
@pytest.mark.parametrize("customer_spend, customer_budget", [(0, 10), (10, 0)])
@pytest.mark.asyncio
async def test_get_end_user_object(customer_spend, customer_budget):
"""
Scenario 1: normal - get_end_user_object returns the cached user
Scenario 2: user over budget - NOTE: budget enforcement now happens in
common_checks() via _check_end_user_budget(), not in get_end_user_object()
This test verifies that get_end_user_object correctly retrieves the end user
from cache. Budget enforcement is tested separately in test_check_end_user_budget().
"""
end_user_id = "my-test-customer"
_budget = LiteLLM_BudgetTable(max_budget=customer_budget)
end_user_obj = LiteLLM_EndUserTable(
user_id=end_user_id,
spend=customer_spend,
litellm_budget_table=_budget,
blocked=False,
)
# UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts
# and breaks get_end_user_object's typed async_get_cache path.
_cache = UserApiKeyCache()
_key = "end_user_id:{}".format(end_user_id)
await _cache.async_set_cache(
key=_key,
value=end_user_obj,
model_type=LiteLLM_EndUserTable,
)
# get_end_user_object only fetches data - it no longer enforces budget
# Budget enforcement happens in common_checks() via _check_end_user_budget()
result = await get_end_user_object(
end_user_id=end_user_id,
prisma_client="RANDOM VALUE", # type: ignore
user_api_key_cache=_cache,
route="/v1/chat/completions",
)
assert result is not None
assert result.user_id == end_user_id
@pytest.mark.parametrize("customer_spend, customer_budget", [(0, 10), (10, 0)])
@pytest.mark.asyncio
async def test_check_end_user_budget(customer_spend, customer_budget):
"""
Test _check_end_user_budget enforcement:
- Scenario 1: customer_spend=0, customer_budget=10 - should pass (under budget)
- Scenario 2: customer_spend=10, customer_budget=0 - should fail (over budget)
Note: Budget enforcement for end users happens in common_checks() via
_check_end_user_budget(), not in get_end_user_object().
"""
from litellm.proxy.auth.auth_checks import check_end_user_budget
_budget = LiteLLM_BudgetTable(max_budget=customer_budget)
end_user_obj = LiteLLM_EndUserTable(
user_id="my-test-customer",
spend=customer_spend,
litellm_budget_table=_budget,
blocked=False,
)
should_exceed = customer_spend > customer_budget
if not should_exceed:
await check_end_user_budget(
end_user_obj=end_user_obj,
route="/v1/chat/completions",
)
return
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await check_end_user_budget(
end_user_obj=end_user_obj,
route="/v1/chat/completions",
)
# Verify the error has correct info
assert exc_info.value.current_cost == customer_spend
assert exc_info.value.max_budget == customer_budget
@pytest.mark.parametrize(
"model, expect_to_work",
[
("openai/gpt-4o-mini", True),
("openai/gpt-4o", False),
],
)
@pytest.mark.asyncio
async def test_can_key_call_model(model, expect_to_work):
"""
If wildcard model + specific model is used, choose the specific model settings
"""
from litellm.proxy.auth.auth_checks import can_key_call_model
from fastapi import HTTPException
llm_model_list = [
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": "test-api-key",
},
"model_info": {
"id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
"db_model": False,
"access_groups": ["public-openai-models"],
},
},
{
"model_name": "openai/gpt-4o",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "test-api-key",
},
"model_info": {
"id": "0cfcd87f2cb12a783a466888d05c6c89df66db23e01cecd75ec0b83aed73c9ad",
"db_model": False,
"access_groups": ["private-openai-models"],
},
},
]
router = litellm.Router(model_list=llm_model_list)
args = {
"model": model,
"llm_model_list": llm_model_list,
"valid_token": UserAPIKeyAuth(
models=["public-openai-models"],
),
"llm_router": router,
}
if expect_to_work:
await can_key_call_model(**args)
else:
with pytest.raises(Exception, match='is not available for this API key') as e:
await can_key_call_model(**args)
print(e)
@pytest.mark.parametrize(
"model, expect_to_work",
[("openai/gpt-4o", False), ("openai/gpt-4o-mini", True)],
)
@pytest.mark.asyncio
async def test_can_team_call_model(model, expect_to_work):
from litellm.proxy.auth.auth_checks import model_in_access_group
from fastapi import HTTPException
llm_model_list = [
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": "test-api-key",
},
"model_info": {
"id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
"db_model": False,
"access_groups": ["public-openai-models"],
},
},
{
"model_name": "openai/gpt-4o",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "test-api-key",
},
"model_info": {
"id": "0cfcd87f2cb12a783a466888d05c6c89df66db23e01cecd75ec0b83aed73c9ad",
"db_model": False,
"access_groups": ["private-openai-models"],
},
},
]
router = litellm.Router(model_list=llm_model_list)
args = {
"model": model,
"team_models": ["public-openai-models"],
"llm_router": router,
}
if expect_to_work:
assert model_in_access_group(**args)
else:
assert not model_in_access_group(**args)
@pytest.mark.parametrize(
"key_models, model, expect_to_work",
[
(["openai/*"], "openai/gpt-4o", True),
(["openai/*"], "openai/gpt-4o-mini", True),
(["openai/*"], "openaiz/gpt-4o-mini", False),
(["bedrock/*"], "bedrock/anthropic.claude-3-5-sonnet-20240620", True),
(["bedrock/*"], "bedrockz/anthropic.claude-3-5-sonnet-20240620", False),
(["bedrock/us.*"], "bedrock/us.amazon.nova-micro-v1:0", True),
(["openai/*"], "ft:gpt-4-0613", True),
(["openai/*"], "bedrockz/ft:gpt-4-0613", False),
],
)
@pytest.mark.asyncio
async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_work):
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.auth_checks import can_key_call_model
llm_model_list = [
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": "test-api-key",
},
"model_info": {
"id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
"db_model": False,
},
},
{
"model_name": "bedrock/*",
"litellm_params": {
"model": "bedrock/*",
"api_key": "test-api-key",
},
"model_info": {
"id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
"db_model": False,
},
},
{
"model_name": "openai/gpt-4o",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "test-api-key",
},
"model_info": {
"id": "0cfcd87f2cb12a783a466888d05c6c89df66db23e01cecd75ec0b83aed73c9ad",
"db_model": False,
},
},
]
router = litellm.Router(model_list=llm_model_list)
user_api_key_object = UserAPIKeyAuth(
models=key_models,
)
if expect_to_work:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=user_api_key_object,
llm_router=router,
)
else:
with pytest.raises(ProxyException):
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=user_api_key_object,
llm_router=router,
)
@pytest.mark.parametrize(
"key_models, model, expect_to_work",
[
# After a cost-map reload, add_known_models() updates anthropic_models so
# the anthropic/* wildcard can match a newly-added Anthropic model.
(["anthropic/*"], "claude-brand-new-model-reload-test", True),
# Wrong provider wildcard must still be denied even after reload.
(["openai/*"], "claude-brand-new-model-reload-test", False),
],
)
@pytest.mark.asyncio
async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_to_work):
"""
Regression test: after a cost-map hot-reload, calling
add_known_models(model_cost_map=new_map) must update litellm.anthropic_models
so that the anthropic/* wildcard correctly grants (or denies) access to
newly-added models.
Root cause: both reload paths in proxy_server.py only updated
litellm.model_cost but never re-ran add_known_models(), so the provider sets
stayed stale and wildcard matching failed for new models.
Fix: each reload now calls litellm.add_known_models(model_cost_map=new_map)
with the fetched map passed explicitly to avoid any reference ambiguity.
"""
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.auth_checks import can_key_call_model
# Build a new cost map that includes the brand-new model — exactly what
# proxy_server.py receives from get_model_cost_map() during a reload.
new_cost_map = dict(litellm.model_cost)
new_cost_map[model] = {
"litellm_provider": "anthropic",
"max_tokens": 8192,
"input_cost_per_token": 0.000003,
"output_cost_per_token": 0.000015,
}
original_model_cost = litellm.model_cost
litellm.model_cost = new_cost_map
# Confirm the model is NOT yet in the provider set before reload propagation.
assert model not in litellm.anthropic_models
# Simulate what proxy_server.py now does after every reload.
litellm.add_known_models(model_cost_map=new_cost_map)
# After add_known_models(), the model must be in the set.
assert model in litellm.anthropic_models
llm_model_list = [
{
"model_name": "anthropic/*",
"litellm_params": {"model": "anthropic/*", "api_key": "test-api-key"},
"model_info": {"id": "test-id-anthropic-wildcard", "db_model": False},
},
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*", "api_key": "test-api-key"},
"model_info": {"id": "test-id-openai-wildcard", "db_model": False},
},
]
router = litellm.Router(model_list=llm_model_list)
user_api_key_object = UserAPIKeyAuth(models=key_models)
try:
if expect_to_work:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=user_api_key_object,
llm_router=router,
)
else:
with pytest.raises(ProxyException):
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=user_api_key_object,
llm_router=router,
)
finally:
litellm.model_cost = original_model_cost
litellm.anthropic_models.discard(model)
@pytest.mark.asyncio
async def test_add_known_models_explicit_map_updates_provider_sets():
"""
Regression test: after a cost-map hot-reload, calling
add_known_models(model_cost_map=new_map) with the new map passed explicitly
must add any new provider models to the correct provider sets so that
wildcard access checks (anthropic/*, openai/*, …) work immediately.
This covers the proxy_server.py fix where both reload paths now call
litellm.add_known_models(model_cost_map=new_model_cost_map) instead of
relying on the module-level model_cost being up to date.
"""
fake_new_model = "claude-brand-new-explicit-map-test"
# Baseline: the model must not be in the sets before we do anything.
assert fake_new_model not in litellm.anthropic_models
new_cost_map = dict(litellm.model_cost)
new_cost_map[fake_new_model] = {
"litellm_provider": "anthropic",
"max_tokens": 8192,
"input_cost_per_token": 0.000003,
"output_cost_per_token": 0.000015,
}
# Simulate what proxy_server.py does on reload.
original_model_cost = litellm.model_cost
litellm.model_cost = new_cost_map
litellm.add_known_models(model_cost_map=new_cost_map)
try:
assert fake_new_model in litellm.anthropic_models, (
"add_known_models(model_cost_map=...) did not add the new model to "
"litellm.anthropic_models — wildcard access checks would fail."
)
finally:
# Clean up: restore original state.
litellm.model_cost = original_model_cost
litellm.anthropic_models.discard(fake_new_model)
@pytest.mark.asyncio
async def test_is_valid_fallback_model():
from litellm.proxy.auth.auth_checks import is_valid_fallback_model
from litellm import Router
router = Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "openai/gpt-3.5-turbo"},
}
]
)
try:
await is_valid_fallback_model(
model="gpt-3.5-turbo", llm_router=router, user_model=None
)
except Exception as e:
pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}")
with pytest.raises(Exception, match="Invalid") as exc_info:
await is_valid_fallback_model(
model="gpt-4o", llm_router=router, user_model=None
)
e = exc_info.value
assert "Invalid" in str(e)
@pytest.mark.parametrize(
"token_spend, max_budget, expect_budget_error",
[
(5.0, 10.0, False), # Under budget
(10.0, 10.0, True), # At budget limit
(15.0, 10.0, True), # Over budget
],
)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_check(
token_spend, max_budget, expect_budget_error
):
"""
Test if virtual key budget checks work as expected:
1. Triggers budget alert for all cases
2. Raises BudgetExceededError when spend >= max_budget
"""
from litellm.proxy.auth.auth_checks import virtual_key_max_budget_check
# Setup test data
valid_token = UserAPIKeyAuth(
token="test-token",
spend=token_spend,
max_budget=max_budget,
user_id="test-user",
key_alias="test-key",
)
user_obj = LiteLLM_UserTable(
user_id="test-user",
user_email="test@email.com",
max_budget=None,
)
proxy_logging_obj = ProxyLogging(
user_api_key_cache=None,
)
# Track if budget alert was called
alert_called = False
async def mock_budget_alert(*args, **kwargs):
nonlocal alert_called
alert_called = True
proxy_logging_obj.budget_alerts = mock_budget_alert
if expect_budget_error:
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
assert exc_info.value.current_cost == token_spend
assert exc_info.value.max_budget == max_budget
else:
await virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
await asyncio.sleep(1)
# Verify budget alert was triggered
assert alert_called, "Budget alert should be triggered"
@pytest.mark.parametrize(
"model, team_models, expect_to_work",
[
("gpt-4", ["gpt-4"], True), # exact match
("gpt-4", ["all-proxy-models"], True), # all-proxy-models access
("gpt-4", ["*"], True), # wildcard access
("gpt-4", ["openai/*"], True), # openai wildcard access
(
"bedrock/anthropic.claude-3-5-sonnet-20240620",
["bedrock/*"],
True,
), # wildcard access
(
"bedrockz/anthropic.claude-3-5-sonnet-20240620",
["bedrock/*"],
False,
), # non-match wildcard access
("bedrock/very_new_model", ["bedrock/*"], True), # bedrock wildcard access
(
"bedrock/claude-3-5-sonnet-20240620",
["bedrock/claude-*"],
True,
), # match on pattern
(
"bedrock/claude-3-6-sonnet-20240620",
["bedrock/claude-3-5-*"],
False,
), # don't match on pattern
("openai/gpt-4o", ["openai/*"], True), # openai wildcard access
("gpt-4", ["gpt-3.5-turbo"], False), # model not in allowed list
("claude-3", [], True), # empty model list (allows all)
],
)
@pytest.mark.asyncio
async def test_can_team_access_model(model, team_models, expect_to_work):
"""
Test cases for can_team_access_model:
1. Exact model match
2. all-proxy-models access
3. Wildcard (*) access
4. OpenAI wildcard access
5. Model not in allowed list
6. Empty model list
7. None model list
"""
try:
team_object = LiteLLM_TeamTable(
team_id="test-team",
models=team_models,
)
result = await can_team_access_model(
model=model,
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
if not expect_to_work:
pytest.fail(
f"Expected model access check to fail for model={model}, team_models={team_models}"
)
except Exception as e:
if expect_to_work:
pytest.fail(
f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {str(e)}"
)
@pytest.mark.parametrize(
"spend, soft_budget, expect_alert",
[
(100, 50, True), # Over soft budget
(50, 50, True), # At soft budget
(25, 50, False), # Under soft budget
(100, None, False), # No soft budget set
],
)
@pytest.mark.asyncio
async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
"""
Test cases for _virtual_key_soft_budget_check:
1. Spend over soft budget
2. Spend at soft budget
3. Spend under soft budget
4. No soft budget set
"""
alert_triggered = False
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered
alert_triggered = True
assert type == "soft_budget"
assert isinstance(user_info, CallInfo)
valid_token = UserAPIKeyAuth(
token="test-token",
spend=spend,
soft_budget=soft_budget,
user_id="test-user",
team_id="test-team",
key_alias="test-key",
)
proxy_logging_obj = MockProxyLogging()
await virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(0.1) # Allow time for the alert task to complete
assert (
alert_triggered == expect_alert
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
@pytest.mark.parametrize(
"spend, soft_budget, expect_alert, metadata, expected_alert_emails",
[
(
100,
50,
False,
None,
None,
), # Over soft budget, no metadata - no alert_emails configured, so no alert
(
50,
50,
False,
None,
None,
), # At soft budget, no metadata - no alert_emails configured, so no alert
(25, 50, False, None, None), # Under soft budget
(100, None, False, None, None), # No soft budget set
(
100,
50,
True,
{"soft_budget_alerting_emails": ["team1@example.com", "team2@example.com"]},
["team1@example.com", "team2@example.com"],
), # Over soft budget with list of emails
(
100,
50,
True,
{"soft_budget_alerting_emails": "team1@example.com,team2@example.com"},
["team1@example.com", "team2@example.com"],
), # Over soft budget with comma-separated emails
(
100,
50,
True,
{
"soft_budget_alerting_emails": [
"team1@example.com",
"",
" ",
"team2@example.com",
]
},
["team1@example.com", "team2@example.com"],
), # Over soft budget with empty strings filtered
],
)
@pytest.mark.asyncio
async def test_team_soft_budget_check(
spend, soft_budget, expect_alert, metadata, expected_alert_emails
):
"""
Test cases for _team_soft_budget_check:
1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
2. Spend at soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
3. Spend under soft budget - should not trigger alert
4. No soft budget set - should not trigger alert
5. Team with alert emails in metadata (list) - should include alert_emails in CallInfo
6. Team with alert emails in metadata (comma-separated string) - should parse and include alert_emails
7. Team with alert emails containing empty strings - should filter them out
"""
alert_triggered = False
captured_call_info = None
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered, captured_call_info
alert_triggered = True
captured_call_info = user_info
assert type == "soft_budget"
assert isinstance(user_info, CallInfo)
valid_token = UserAPIKeyAuth(
token="test-token",
user_id="test-user",
team_id="test-team",
team_alias="test-team-alias",
key_alias="test-key",
)
team_object = LiteLLM_TeamTable(
team_id="test-team",
spend=spend,
soft_budget=soft_budget,
max_budget=100.0,
metadata=metadata,
)
proxy_logging_obj = MockProxyLogging()
await _team_soft_budget_check(
team_object=team_object,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(0.1) # Allow time for the alert task to complete
assert (
alert_triggered == expect_alert
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
if expect_alert:
assert captured_call_info is not None
assert captured_call_info.team_id == "test-team"
assert captured_call_info.spend == spend
assert captured_call_info.soft_budget == soft_budget
assert captured_call_info.event_group == Litellm_EntityType.TEAM
# Verify alert_emails if expected
if expected_alert_emails is not None:
assert captured_call_info.alert_emails == expected_alert_emails
else:
assert (
captured_call_info.alert_emails is None
or captured_call_info.alert_emails == []
)
@pytest.mark.asyncio
async def test_can_user_call_model():
from litellm.proxy.auth.auth_checks import can_user_call_model
from litellm.proxy._types import ProxyException
from litellm import Router
router = Router(
model_list=[
{
"model_name": "anthropic-claude",
"litellm_params": {"model": "anthropic/anthropic-claude"},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "test-api-key"},
},
]
)
args = {
"model": "anthropic-claude",
"llm_router": router,
"user_object": LiteLLM_UserTable(
user_id="testuser21@mycompany.com",
max_budget=None,
spend=0.0042295,
model_max_budget={},
model_spend={},
user_email="testuser@mycompany.com",
models=["gpt-3.5-turbo"],
),
}
with pytest.raises(ProxyException) as e:
await can_user_call_model(**args)
args["model"] = "gpt-3.5-turbo"
await can_user_call_model(**args)
@pytest.mark.asyncio
async def test_can_user_call_model_with_no_default_models():
from litellm.proxy.auth.auth_checks import can_user_call_model
from litellm.proxy._types import ProxyException, SpecialModelNames
from unittest.mock import MagicMock
args = {
"model": "anthropic-claude",
"llm_router": MagicMock(),
"user_object": LiteLLM_UserTable(
user_id="testuser21@mycompany.com",
max_budget=None,
spend=0.0042295,
model_max_budget={},
model_spend={},
user_email="testuser@mycompany.com",
models=[SpecialModelNames.no_default_models.value],
),
}
with pytest.raises(ProxyException) as e:
await can_user_call_model(**args)
@pytest.mark.asyncio
async def test_get_fuzzy_user_object():
from litellm.proxy.auth.auth_checks import _get_fuzzy_user_object
from unittest.mock import AsyncMock, MagicMock
# Setup mock Prisma client
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
mock_prisma.db.litellm_usertable = MagicMock()
# Mock user data
test_user = LiteLLM_UserTable(
user_id="test_123",
sso_user_id="sso_123",
user_email="test@example.com",
organization_memberships=[],
max_budget=None,
)
# Test 1: Find user by SSO ID
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma, sso_user_id="sso_123", user_email="test@example.com"
)
assert result == test_user
mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
)
# Test 2: SSO ID not found, find by email
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
mock_prisma.db.litellm_usertable.update = AsyncMock()
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma,
sso_user_id="new_sso_456",
user_email="test@example.com",
)
assert result == test_user
mock_prisma.db.litellm_usertable.find_first.assert_called_with(
where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
include={"organization_memberships": True},
)
# Test 3: Verify background SSO update task when user found by email
await asyncio.sleep(0.1) # Allow time for background task
mock_prisma.db.litellm_usertable.update.assert_called_with(
where={"user_id": "test_123"}, data={"sso_user_id": "new_sso_456"}
)
# Test 4: User not found by either method
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=None)
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma,
sso_user_id="unknown_sso",
user_email="unknown@example.com",
)
assert result is None
# Test 5: Only email provided (no SSO ID)
mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma, user_email="test@example.com"
)
assert result == test_user
mock_prisma.db.litellm_usertable.find_first.assert_called_with(
where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
include={"organization_memberships": True},
)
# Test 6: Only SSO ID provided (no email)
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma, sso_user_id="sso_123"
)
assert result == test_user
mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
)
@pytest.mark.parametrize(
"model, alias_map, expect_to_work",
[
("gpt-4", {"gpt-4": "gpt-4-team1"}, True), # model matches alias value
("gpt-5", {"gpt-4": "gpt-4-team1"}, False),
],
)
@pytest.mark.asyncio
async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work):
"""
Test if can_key_call_model correctly handles model aliases in the token
"""
from litellm.proxy.auth.auth_checks import can_key_call_model
llm_model_list = [
{
"model_name": "gpt-4-team1",
"litellm_params": {
"model": "gpt-4",
"api_key": "test-api-key",
},
}
]
router = litellm.Router(model_list=llm_model_list)
user_api_key_object = UserAPIKeyAuth(
models=[
"gpt-4-team1",
],
team_model_aliases=alias_map,
)
if expect_to_work:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=user_api_key_object,
llm_router=router,
)
else:
with pytest.raises(Exception, match='is not available for this API key') as e:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=user_api_key_object,
llm_router=router,
)
# ---------------------------------------------------------------------------
# Access group cache helpers (_cache_access_object, _delete_cache_access_object)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_cache_access_object():
"""Test _cache_access_object stores access group in cache with correct key."""
from litellm.proxy.auth.auth_checks import cache_access_object
from litellm.proxy._types import LiteLLM_AccessGroupTable
cache = DualCache()
ag_id = "ag-test-123"
ag_table = LiteLLM_AccessGroupTable(
access_group_id=ag_id,
access_group_name="test-group",
access_model_names=["gpt-4"],
)
await cache_access_object(
access_group_id=ag_id,
access_group_table=ag_table,
user_api_key_cache=cache,
)
cached = await cache.async_get_cache(key=f"access_group_id:{ag_id}")
assert cached is not None
if isinstance(cached, dict):
assert cached.get("access_group_id") == ag_id
assert cached.get("access_group_name") == "test-group"
else:
assert cached.access_group_id == ag_id
assert cached.access_group_name == "test-group"
@pytest.mark.asyncio
async def test_delete_cache_access_object():
"""Test _delete_cache_access_object removes access group from in-memory cache."""
from litellm.proxy.auth.auth_checks import delete_cache_access_object
from litellm.proxy._types import LiteLLM_AccessGroupTable
cache = DualCache()
ag_id = "ag-delete-test"
ag_table = LiteLLM_AccessGroupTable(
access_group_id=ag_id,
access_group_name="to-delete",
)
await cache.async_set_cache(key=f"access_group_id:{ag_id}", value=ag_table, ttl=60)
await delete_cache_access_object(access_group_id=ag_id, user_api_key_cache=cache)
cached = await cache.async_get_cache(key=f"access_group_id:{ag_id}")
assert cached is None
# ---------------------------------------------------------------------------
# Access group resource fetchers (_get_models_from_access_groups, _get_agent_ids_from_access_groups)
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"resource_field, access_group_data, expected",
[
(
"access_model_names",
{"access_group_id": "ag-1", "access_model_names": ["gpt-4", "claude-3"]},
["gpt-4", "claude-3"],
),
(
"access_agent_ids",
{"access_group_id": "ag-2", "access_agent_ids": ["agent-a", "agent-b"]},
["agent-a", "agent-b"],
),
(
"access_model_names",
{"access_group_id": "ag-3", "access_model_names": []},
[],
),
],
)
@pytest.mark.asyncio
async def test_get_resources_from_access_groups(
resource_field, access_group_data, expected
):
"""Test _get_resources_from_access_groups returns correct resource list from access groups."""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import LiteLLM_AccessGroupTable
from litellm.proxy.auth.auth_checks import (
get_agent_ids_from_access_groups,
get_models_from_access_groups,
)
ag_table = LiteLLM_AccessGroupTable(
access_group_id=access_group_data["access_group_id"],
access_group_name="test",
access_model_names=access_group_data.get("access_model_names", []),
access_agent_ids=access_group_data.get("access_agent_ids", []),
)
with patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
return_value=ag_table,
):
if resource_field == "access_model_names":
result = await get_models_from_access_groups(
access_group_ids=[access_group_data["access_group_id"]],
prisma_client=MagicMock(),
user_api_key_cache=DualCache(),
)
else:
result = await get_agent_ids_from_access_groups(
access_group_ids=[access_group_data["access_group_id"]],
prisma_client=MagicMock(),
user_api_key_cache=DualCache(),
)
assert sorted(result) == sorted(expected)
@pytest.mark.asyncio
async def test_get_models_from_access_groups_empty_ids():
"""Test _get_models_from_access_groups returns empty list when access_group_ids is empty."""
from litellm.proxy.auth.auth_checks import get_models_from_access_groups
result = await get_models_from_access_groups(access_group_ids=[])
assert result == []
# ---------------------------------------------------------------------------
# can_team_access_model with access_group_ids fallback
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_can_team_access_model_via_access_group_ids():
"""Test can_team_access_model allows access when team has access_group_ids granting model access."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import can_team_access_model
team_object = LiteLLM_TeamTable(
team_id="test-team",
models=[],
access_group_ids=["ag-with-gpt4"],
)
with patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["gpt-4"],
):
result = await can_team_access_model(
model="gpt-4",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
assert result is True
@pytest.mark.asyncio
async def test_can_team_access_model_access_group_ids_denied():
"""Test can_team_access_model denies when neither team models nor access_group_ids grant access."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy._types import ProxyException
team_object = LiteLLM_TeamTable(
team_id="test-team",
models=["gpt-3.5-turbo"],
access_group_ids=["ag-other"],
)
with patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["claude-3"],
):
with pytest.raises(ProxyException):
await can_team_access_model(
model="gpt-4",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
# ---------------------------------------------------------------------------
# can_key_call_model with access_group_ids fallback
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_can_key_call_model_via_access_group_ids():
"""Test can_key_call_model allows access when key has access_group_ids granting model access."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import can_key_call_model
user_api_key_object = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=["ag-with-gpt4"],
)
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "openai/gpt-4", "api_key": "test"},
}
]
)
with patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["gpt-4"],
):
await can_key_call_model(
model="gpt-4",
llm_model_list=[],
valid_token=user_api_key_object,
llm_router=router,
)
# ---------------------------------------------------------------------------
# _key_access_group_grants_model (key access group overriding team restriction)
# ---------------------------------------------------------------------------
def _patch_proxy_server_globals():
"""Patch proxy_server's prisma_client and user_api_key_cache to non-None mocks
so the helper's None-guard doesn't short-circuit. The actual values don't
matter because get_access_object is patched separately to return fixtures."""
from unittest.mock import MagicMock, patch
return [
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
]
def _fake_access_group(
access_group_id: str,
access_model_names=None,
assigned_team_ids=None,
assigned_key_ids=None,
):
from litellm.proxy._types import LiteLLM_AccessGroupTable
return LiteLLM_AccessGroupTable(
access_group_id=access_group_id,
access_group_name=access_group_id,
access_model_names=access_model_names or [],
assigned_team_ids=assigned_team_ids or [],
assigned_key_ids=assigned_key_ids or [],
)
@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_team_authorized():
"""Group's assigned_team_ids includes the key's team and grants the model → True.
This is the happy path equivalent of Andres's report: admin creates an
access group with assigned_team_ids=[team-a], grants claude-haiku-4-5,
attaches it to a key on team-a. Override fires.
"""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import key_access_group_grants_model
valid_token = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=["premium-group"],
team_id="team-a",
)
team_object = LiteLLM_TeamTable(
team_id="team-a",
models=["mock-success"],
access_group_ids=[], # deliberately not synced — the access group itself authorizes
)
fake_ag = _fake_access_group(
access_group_id="premium-group",
access_model_names=["claude-haiku-4-5"],
assigned_team_ids=["team-a"],
)
patches = _patch_proxy_server_globals() + [
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
return_value=fake_ag,
),
]
for p in patches:
p.start()
try:
assert (
await key_access_group_grants_model(
model="claude-haiku-4-5",
valid_token=valid_token,
team_object=team_object,
llm_router=None,
)
is True
)
finally:
for p in patches:
p.stop()
@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_key_directly_authorized():
"""Group's assigned_key_ids includes the key's token and grants the model → True.
Per-key authorization path: an admin scopes a group directly to a key
(assigned_key_ids) without listing the team.
"""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import key_access_group_grants_model
valid_token = UserAPIKeyAuth(
token="test-token-hashed",
models=[],
access_group_ids=["per-key-group"],
team_id="team-a",
)
team_object = LiteLLM_TeamTable(
team_id="team-a",
models=["mock-success"],
access_group_ids=[],
)
fake_ag = _fake_access_group(
access_group_id="per-key-group",
access_model_names=["claude-haiku-4-5"],
assigned_team_ids=[],
assigned_key_ids=["test-token-hashed"],
)
patches = _patch_proxy_server_globals() + [
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
return_value=fake_ag,
),
]
for p in patches:
p.start()
try:
assert (
await key_access_group_grants_model(
model="claude-haiku-4-5",
valid_token=valid_token,
team_object=team_object,
llm_router=None,
)
is True
)
finally:
for p in patches:
p.stop()
@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_key_has_no_groups():
"""Key with no access_group_ids → False (early return, no DB read)."""
from litellm.proxy.auth.auth_checks import key_access_group_grants_model
valid_token = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=[],
team_id="team-a",
)
team_object = LiteLLM_TeamTable(
team_id="team-a",
models=["mock-success"],
access_group_ids=["any-group"],
)
assert (
await key_access_group_grants_model(
model="claude-haiku-4-5",
valid_token=valid_token,
team_object=team_object,
llm_router=None,
)
is False
)
@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_group_does_not_cover_model():
"""Group authorizes the team but does not grant the requested model → False."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import key_access_group_grants_model
valid_token = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=["basic-group"],
team_id="team-a",
)
team_object = LiteLLM_TeamTable(
team_id="team-a",
models=["mock-success"],
access_group_ids=[],
)
fake_ag = _fake_access_group(
access_group_id="basic-group",
access_model_names=["gpt-4o-mini"],
assigned_team_ids=["team-a"],
)
patches = _patch_proxy_server_globals() + [
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
return_value=fake_ag,
),
]
for p in patches:
p.start()
try:
assert (
await key_access_group_grants_model(
model="claude-haiku-4-5",
valid_token=valid_token,
team_object=team_object,
llm_router=None,
)
is False
)
finally:
for p in patches:
p.stop()
@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_group_authorizes_neither():
"""
Bypass regression test: a team member sets a foreign access group on their
key. The group grants the requested model but its assigned_team_ids /
assigned_key_ids do not include this caller's team or token. Override is
denied — the team's 401 propagates.
"""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import key_access_group_grants_model
valid_token = UserAPIKeyAuth(
token="team-a-token",
models=[],
access_group_ids=["team-b-premium"],
team_id="team-a",
)
team_object = LiteLLM_TeamTable(
team_id="team-a",
models=["mock-success"],
access_group_ids=[],
)
fake_ag = _fake_access_group(
access_group_id="team-b-premium",
access_model_names=["claude-opus-4-5"],
assigned_team_ids=["team-b"],
assigned_key_ids=["team-b-token"],
)
patches = _patch_proxy_server_globals() + [
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
return_value=fake_ag,
),
]
for p in patches:
p.start()
try:
assert (
await key_access_group_grants_model(
model="claude-opus-4-5",
valid_token=valid_token,
team_object=team_object,
llm_router=None,
)
is False
)
finally:
for p in patches:
p.stop()
@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_get_access_object_raises():
"""Group lookup failure (404, network, etc.) is treated as no authorization."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import key_access_group_grants_model
valid_token = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=["missing-group"],
team_id="team-a",
)
team_object = LiteLLM_TeamTable(
team_id="team-a",
models=["mock-success"],
access_group_ids=[],
)
patches = _patch_proxy_server_globals() + [
patch(
"litellm.proxy.auth.auth_checks.get_access_object",
new_callable=AsyncMock,
side_effect=Exception("not found"),
),
]
for p in patches:
p.start()
try:
assert (
await key_access_group_grants_model(
model="claude-haiku-4-5",
valid_token=valid_token,
team_object=team_object,
llm_router=None,
)
is False
)
finally:
for p in patches:
p.stop()
@pytest.fixture
def db_success_hook() -> Iterator[AsyncMock]:
hook: Final = AsyncMock()
with patch(
"litellm.proxy.proxy_server.proxy_logging_obj",
MagicMock(service_logging_obj=MagicMock(async_service_success_hook=hook)),
):
yield hook
async def _db_service_call_types(hook: AsyncMock) -> tuple[str, ...]:
await asyncio.sleep(0)
return tuple(call.kwargs["call_type"] for call in hook.await_args_list if call.kwargs["service"] == ServiceTypes.DB)
def _prisma_client_serving(user_id: str) -> SimpleNamespace:
row: Final = {
"user_id": user_id,
"user_role": "internal_user",
"teams": [],
"spend": 0.0,
"models": [],
"metadata": "{}",
"allowed_cache_controls": [],
"policies": [],
"model_spend": "{}",
"model_max_budget": "{}",
"organization_memberships": [],
}
engine: Final = SimpleNamespace(query=AsyncMock(return_value={"data": {"result": row}}), stop=lambda: None)
generated_client: Final = Prisma()
generated_client._engine = engine
return SimpleNamespace(db=PrismaWrapper(original_prisma=generated_client, iam_token_db_auth=False))
@pytest.mark.asyncio
async def test_get_user_object_cache_hit_emits_no_postgres_service_event(db_success_hook: AsyncMock) -> None:
user_id: Final = f"cached-user-{uuid.uuid4()}"
cache: Final = UserApiKeyCache()
await cache.async_set_cache(key=user_id, value=LiteLLM_UserTable(user_id=user_id, user_role="internal_user"))
result: Final = await get_user_object(
user_id=user_id,
prisma_client=MagicMock(),
user_api_key_cache=cache,
user_id_upsert=False,
parent_otel_span="auth-span",
)
assert result is not None and result.user_id == user_id
assert await _db_service_call_types(db_success_hook) == ()
@pytest.mark.asyncio
async def test_get_org_object_cache_hit_emits_no_postgres_service_event(db_success_hook: AsyncMock) -> None:
org_id: Final = f"cached-org-{uuid.uuid4()}"
cache: Final = UserApiKeyCache()
await cache.async_set_cache(
key=f"org_id:{org_id}",
value=LiteLLM_OrganizationTable(
organization_id=org_id, budget_id="b", models=[], created_by="t", updated_by="t"
),
)
result: Final = await get_org_object(
org_id=org_id,
prisma_client=MagicMock(),
user_api_key_cache=cache,
parent_otel_span="auth-span",
)
assert result is not None and result.organization_id == org_id
assert await _db_service_call_types(db_success_hook) == ()
@pytest.mark.asyncio
async def test_get_user_object_cache_miss_emits_exactly_one_postgres_get_user_object_event(
db_success_hook: AsyncMock,
) -> None:
user_id: Final = f"db-user-{uuid.uuid4()}"
prisma_client: Final = _prisma_client_serving(user_id)
result: Final = await get_user_object(
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=UserApiKeyCache(),
user_id_upsert=False,
parent_otel_span="auth-span",
)
assert result is not None and result.user_id == user_id
assert await _db_service_call_types(db_success_hook) == ("get_user_object",)
assert db_success_hook.await_args_list[0].kwargs["parent_otel_span"] == "auth-span"
@pytest.mark.parametrize("entry_first", [True, False])
def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_routes_to(
monkeypatch: pytest.MonkeyPatch, entry_first: bool
) -> None:
"""chain-entry resolves one hop to local-free and is served by local-free's own free
deployment, so an over-budget key is waived for it; local-free by name resolves to paid-gpt
and stays enforced. local-free's alias is a hop the router never takes for chain-entry, and
the per-name verdict cache must not let either name's verdict leak into the other's."""
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
router: Final = litellm.Router(
model_list=[
{
"model_name": "local-free",
"litellm_params": {
"model": "ollama/qwen3:0.6b",
"api_base": "http://localhost:11434",
"input_cost_per_token": 0,
"output_cost_per_token": 0,
},
},
{
"model_name": "paid-gpt",
"litellm_params": {
"model": "gpt-4o",
"api_key": "fake",
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
},
},
],
model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"},
)
expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False}
order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order}
assert verdicts == expected
assert {name: is_model_cost_zero(model=name, llm_router=router) for name in order} == expected
MCP_DISCOVERY_REST_ROUTES = [
"/mcp-rest/tools/list",
"/v1/mcp/tools",
]
@pytest.mark.parametrize("route", MCP_DISCOVERY_REST_ROUTES)
@pytest.mark.asyncio
async def test_mcp_tool_discovery_route_bypasses_team_budget(route):
from litellm.proxy.auth.auth_checks import common_checks
team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
result = await common_checks(
request_body={},
team_object=team_object,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=AsyncMock(),
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
request=MagicMock(),
)
assert result is True
@pytest.mark.parametrize("method", ["initialize", "notifications/initialized", "ping", "tools/list"])
@pytest.mark.parametrize("route", ["/mcp", "/mcp/some-server"])
@pytest.mark.asyncio
async def test_mcp_jsonrpc_zero_spend_method_bypasses_team_budget(route, method):
from litellm.proxy.auth.auth_checks import common_checks
team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
result = await common_checks(
request_body={"jsonrpc": "2.0", "id": 1, "method": method, "params": {}},
team_object=team_object,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=AsyncMock(),
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
request=MagicMock(),
)
assert result is True
@pytest.mark.parametrize(
"route,request_body",
[
("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
("/mcp/some-server", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
("/mcp-rest/tools/call", {"name": "add", "arguments": {}}),
("/mcp/tools/call", {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}),
("/mcp/tools", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
("/mcp/tools/list", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
],
)
@pytest.mark.asyncio
async def test_mcp_tool_call_route_still_enforces_team_budget(route, request_body):
from litellm.proxy.auth.auth_checks import common_checks
team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
with pytest.raises(litellm.BudgetExceededError):
await common_checks(
request_body=request_body,
team_object=team_object,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=AsyncMock(),
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
request=MagicMock(),
)
@pytest.mark.parametrize(
"route,request_body,expected",
[
("/mcp-rest/tools/list", {}, True),
("/v1/mcp/tools", {}, True),
("/mcp/tools", {}, False),
("/mcp/tools/list", {}, False),
("/mcp/tools", {"method": "tools/call"}, False),
("/mcp/tools", {"method": "tools/list"}, True),
("/mcp", {"method": "initialize"}, True),
("/mcp", {"method": "notifications/initialized"}, True),
("/mcp", {"method": "ping"}, True),
("/mcp", {"method": "tools/list"}, True),
("/mcp/some-server", {"method": "tools/list"}, True),
("/mcp", {"method": "tools/call"}, False),
("/mcp/some-server", {"method": "tools/call"}, False),
("/mcp", {}, False),
("/mcp", {"method": ["tools/list"]}, False),
("/mcp", {"method": {"name": "ping"}}, False),
("/mcp-rest/tools/call", {"method": "initialize"}, False),
("/mcp/tools/call", {"method": "initialize"}, False),
("/v1/chat/completions", {"method": "initialize"}, False),
("/mcp-rest/other", {"method": "tools/list"}, False),
],
)
def test_is_mcp_discovery_request(route, request_body, expected):
from litellm.proxy.auth.auth_checks import is_mcp_discovery_request
assert is_mcp_discovery_request(route=route, request_body=request_body) is expected
@pytest.mark.asyncio
async def test_get_skill_names_from_access_groups_reads_cached_group_rows():
from litellm.proxy._types import LiteLLM_AccessGroupTable
from litellm.proxy.auth.auth_checks import cache_access_object, get_skill_names_from_access_groups
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache = UserApiKeyCache()
for group_id, skills in (("group-a", ["skill-a", "shared"]), ("group-b", ["skill-b", "shared"])):
await cache_access_object(
access_group_id=group_id,
access_group_table=LiteLLM_AccessGroupTable(
access_group_id=group_id, access_group_name=group_id, access_skill_names=skills
),
user_api_key_cache=cache,
)
untouched_prisma = MagicMock()
untouched_prisma.db.litellm_accessgrouptable.find_unique = AsyncMock(
side_effect=AssertionError("cache hit must not reach the DB")
)
assert await get_skill_names_from_access_groups(
access_group_ids=("group-a", "group-b"),
prisma_client=untouched_prisma,
user_api_key_cache=cache,
proxy_logging_obj=MagicMock(),
) == frozenset({"skill-a", "skill-b", "shared"})
class _Pipeline:
def __init__(self, send: Callable[[Sequence[str]], Awaitable[None]]) -> None:
self._send = send
self._messages: list[str] = []
def publish(self, channel: str, message: str) -> "_Pipeline":
self._messages.append(message)
return self
async def execute(self) -> None:
await self._send(tuple(self._messages))
class _SlowPublishRedisClient:
def __init__(self) -> None:
self.published: list[str] = []
async def _send(self, messages: Sequence[str]) -> None:
await asyncio.sleep(0.001)
self.published.extend(messages)
async def publish(self, channel: str, message: str) -> int:
await self._send((message,))
return 1
def pipeline(self, transaction: bool = True) -> _Pipeline:
return _Pipeline(self._send)
class _WedgedPublishRedisClient:
def __init__(self) -> None:
self.release = asyncio.Event()
async def _send(self, messages: Sequence[str]) -> None:
await self.release.wait()
async def publish(self, channel: str, message: str) -> int:
await self._send((message,))
return 1
def pipeline(self, transaction: bool = True) -> _Pipeline:
return _Pipeline(self._send)
class _PubSubRedisCache:
namespace = None
def __init__(self, client: _SlowPublishRedisClient | _WedgedPublishRedisClient) -> None:
self._client = client
def init_pubsub_client(self) -> _SlowPublishRedisClient | _WedgedPublishRedisClient:
return self._client
async def _drain_publishes() -> None:
await asyncio.gather(*pubsub_module._pending_publishes) # pyright: ignore[reportPrivateUsage] # drain module-level tasks
@pytest.mark.asyncio
async def test_evicting_more_keys_than_the_publisher_holds_pending_broadcasts_every_one(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
client = _SlowPublishRedisClient()
hashed_tokens: Final = tuple(f"token-{index}" for index in range(pubsub_module._MAX_PENDING_PUBLISHES + 476)) # pyright: ignore[reportPrivateUsage] # the cap under test
with patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_PubSubRedisCache(client),
):
await delete_cache_key_objects(
hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=None
)
await _drain_publishes()
assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens)
@pytest.mark.asyncio
async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
hashed_tokens: Final = tuple(
hashlib.sha256(f"token-{index}".encode()).hexdigest()
for index in range(600)
)
for hashed_token in hashed_tokens:
cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})
seeded: Final = tuple(
token for token in hashed_tokens if cache.in_memory_cache_for(token).get_cache(token) is not None
)
with patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_PubSubRedisCache(client),
):
eviction: Final = asyncio.create_task(
delete_cache_key_objects(hashed_tokens=hashed_tokens, user_api_key_cache=cache, proxy_logging_obj=None)
)
await asyncio.sleep(0.05)
still_cached: Final = tuple(
token for token in seeded if cache.in_memory_cache_for(token).get_cache(token) is not None
)
client.release.set()
await eviction
await _drain_publishes()
assert seeded == hashed_tokens
assert still_cached == ()
class _AuthRedisCache:
namespace = None
def __init__(self, entries: dict[str, object]) -> None:
self.entries = dict(entries)
async def async_get_cache(self, key: str, **kwargs: object) -> object | None:
return self.entries.get(key)
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
self.entries[key] = value
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
self.entries[key] = value
def delete_cache(self, key: str) -> None:
self.entries.pop(key, None)
async def async_delete_cache(self, key: str) -> None:
self.entries.pop(key, None)
async def delete_cache_keys(self, keys: Sequence[str]) -> None:
for key in keys:
self.entries.pop(key, None)
@pytest.mark.asyncio
async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_broadcast(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
hashed_tokens: Final = tuple(
hashlib.sha256(f"token-{index}".encode()).hexdigest()
for index in range(600)
)
redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens})
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer
for hashed_token in hashed_tokens:
cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})
with patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_PubSubRedisCache(client),
):
eviction: Final = asyncio.create_task(
delete_cache_key_objects(hashed_tokens=hashed_tokens, user_api_key_cache=cache, proxy_logging_obj=None)
)
await asyncio.sleep(0.05)
refilled: Final = [token for token in hashed_tokens if await cache.async_get_cache(key=token) is not None]
client.release.set()
await eviction
await _drain_publishes()
assert refilled == []
assert redis.entries == {}
Reference in a new issue View git blame Copy permalink