scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(rate-limiting): reserve one concurrency unit per non-racing batch branch, release one per terminal event

Browse source 
Router.abatch_completion's comma-separated model fans a single admission
into several real, independent LLM calls below this hook entirely (it never
re-runs async_pre_call_hook per branch the way model_based_tag_rate_limits_hook's
per-Router-hop admission does). A single reservation for the whole batch was
released by whichever branch finished first, letting the still-running
siblings push real concurrent calls past the configured cap. Confirmed live
against the real Router.abatch_completion: admission fires exactly once per
dispatch, and every branch reliably fires its own terminal event (no
cancellation involved), unlike abatch_completion_fastest_response, whose
cancelled losing branches never fire one at all.

Admission now reserves one unit per comma-separated model for that specific,
non-racing call shape (_non_racing_batch_width), and each branch's own
success/failure event releases just its own share (_release_one_pending_for_call_id),
leaving the existing release-everything path for disconnect/exhausted-fallback
unchanged, since those abort every not-yet-completed branch together with no
per-branch event to rely on. The racing fastest_response variant keeps
reserving a single unit, since reserving per-model there would leak every
cancelled loser's share until the safety TTL.

The sibling-slot-release bug fixed on the parent branch does not reproduce
here: this hook's own admission point never runs more than once per call_id,
so there is never more than one reservation to disambiguate between.
This commit is contained in:
Deepanshu 2026-08-31 12:34:57 -04:00
parent 63299d9df6
commit a0a92093f6
2 changed files with 423 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
litellm/proxy/hooks/global_tag_rate_limits_hook.py
View file

@ -32,6 +32,22 @@ same inherited context (an LLM-judge guardrail, for example) gets its own
isolated entry instead of releasing the outer call's still-pending
reservation early. Confirmed live: a real streaming client disconnect
correctly releases its concurrency reservation through this mechanism.
`Router.abatch_completion`'s comma-separated `model` fans this one admission
out into several real, independent LLM calls below this hook entirely (it
never re-runs `async_pre_call_hook` per branch the way
`model_based_tag_rate_limits_hook`'s per-Router-hop admission does), each of
which reliably fires its own terminal success/failure event -- confirmed
live against the real dispatch. A single concurrency reservation for the
whole batch would be released by whichever branch finishes first, letting
the still-running siblings push real concurrent calls past the configured
cap, so admission instead reserves one unit per comma-separated model (see
`_non_racing_batch_width`) and each branch's own event releases just its own
share (see `_release_one_pending_for_call_id`). The racing
`abatch_completion_fastest_response` variant is excluded from that: it
cancels every losing branch without ever firing a terminal event for it
(confirmed live), so reserving more than the single unit it already does
would leak every losing branch's share until the safety TTL.
"""
import asyncio
@ -142,6 +158,34 @@ def _entry_applies_any_admitted_model(
return any(_entry_applies(entry, tags, key_alias, model) for model in admitted_models)
def _non_racing_batch_width(data: Mapping[str, object], call_type: str) -> int:
"""`Router.abatch_completion` (not `abatch_completion_fastest_response`)
fans this one admission out into one independent real LLM call per
comma-separated model in `model`, and every one of those branches
reliably fires its own terminal success/failure event -- confirmed
empirically, no cancellation involved, unlike the racing
`fastest_response` variant, which cancels every losing branch without
ever firing a terminal event for it and must keep reserving a single
unit released by whichever branch finishes first.
A single concurrency reservation for the whole non-racing dispatch would
get released by whichever branch finishes first, letting the
still-running siblings push real concurrent calls past the configured
cap. Reserving one unit per branch instead, released one at a time as
each branch's own event fires, keeps the count accurate.
Mirrors the exact condition `route_llm_request.py` uses to route into
`abatch_completion` in the first place, so this only ever fires for a
request that will actually take that path.
"""
if call_type != "acompletion" or data.get("fastest_response"):
return 1
model_field: Final = data.get("model")
if not isinstance(model_field, str) or "," not in model_field:
return 1
return len(model_field.split(","))
def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str:
"""Namespaced under `tag_rl:global:` so it never collides with
`model_based_tag_rate_limits_hook`'s own `tag_rl:{model_group}:...` keys."""
@ -195,6 +239,11 @@ class _GlobalTagRateLimitStash:
# an earlier attempt must still get its accounting at success time even
# though the request ultimately serves from a later attempt's model.
admitted_models: frozenset[str] = field(default_factory=frozenset)
# One entry per reserved unit, not one entry per distinct key: a
# non-racing batch dispatch (see _non_racing_batch_width) reserves and
# appends `batch_width` entries for the same key, and each branch's own
# terminal event pops exactly one of them -- see
# _release_one_pending_for_call_id.
pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue
# Keys already charged for this call_id, so a fallback retry (same
# litellm_call_id, different model) renews instead of double-charging.
@ -587,6 +636,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
already_reserved_concurrency_keys: Final = frozenset(
key for key, _partition_key in stash.pending_concurrency_keys
)
non_racing_batch_width: Final = _non_racing_batch_width(data, call_type)
failing_index, values = await self._atomic_check_and_increment(
tuple(
(
@ -596,13 +646,16 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
# A key already charged/reserved for this call_id (an
# earlier fallback attempt for the same request) renews
# at zero net cost instead of charging a second unit.
# Otherwise a concurrency check reserves one unit per
# non-racing batch branch (see _non_racing_batch_width),
# not just one for the whole dispatch.
0.0
if renewal_allowed
and (
(check.unit == "requests" and check.key in stash.charged_request_keys)
or (check.unit == "concurrency" and check.key in already_reserved_concurrency_keys)
)
else 1.0,
else (float(non_racing_batch_width) if check.unit == "concurrency" else 1.0),
self._ttl_for(check.unit, check.entry),
check.unit == "concurrency",
)
@ -617,11 +670,14 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
# Exclude already_reserved_concurrency_keys: that key renewed at
# zero cost above, so re-adding it would make release decrement
# twice for a counter only ever incremented once.
# twice for a counter only ever incremented once. One entry per
# reserved unit (non_racing_batch_width of them, ordinarily 1) --
# see pending_concurrency_keys's own docstring for why.
concurrency_reservations: Final = tuple(
(check.key, _partition_key(check.entry))
for check in atomic_checks
if check.unit == "concurrency" and check.key not in already_reserved_concurrency_keys
for _ in range(non_racing_batch_width)
)
if concurrency_reservations:
stash.pending_concurrency_keys.extend(concurrency_reservations) # mutable-ok: see field's own docstring
@ -645,6 +701,14 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
return data
async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None:
"""Releases every reservation still pending for this call_id at once.
Correct for a disconnect or a chain-exhausted failure: either aborts
every not-yet-completed branch of a non-racing batch together (see
_non_racing_batch_width), so none of them will ever fire its own
terminal event to release its own share individually -- whatever's
still pending here is exactly what those abandoned branches reserved,
no more (any branch that already completed already popped its own
entry via _release_one_pending_for_call_id) and no less."""
stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs))
if stash is None or not stash.pending_concurrency_keys:
return
@ -652,6 +716,21 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
stash.pending_concurrency_keys.clear()
await self._release_keys(release_keys)
async def _release_one_pending_for_call_id(
self, request_kwargs: Mapping[str, object]
) -> tuple[str, _PartitionKey] | None:
"""Releases exactly one reservation, not every reservation currently
pending: a non-racing batch dispatch reserves one unit per
comma-separated model (see _non_racing_batch_width), and each
branch's own terminal event must only release its own share, not a
still-running sibling's. Entries under one call_id are otherwise
fungible (same key repeated), so which single entry gets popped
doesn't matter."""
stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs))
if stash is None or not stash.pending_concurrency_keys:
return None
return stash.pending_concurrency_keys.pop() # mutable-ok: see field's own docstring
async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None:
await self._release_pending_for_call_id(request_data)
@ -683,11 +762,15 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
end_time: datetime | None,
) -> None:
# Always release regardless of which hook raised: this hook's own
# rejection never reserves a slot, so pending_concurrency_keys is
# already empty in that case and the check below no-ops; a rejection
# from model_based_tag_rate_limits_hook (same error marker) can still
# land after this hook already reserved its own slot.
await self._release_pending_for_call_id(kwargs)
# rejection never reserves a slot, so there is nothing to pop in
# that case; a rejection from model_based_tag_rate_limits_hook (same
# error marker) can still land after this hook already reserved its
# own slot. Only this one branch's own share, not every reservation
# still pending for a non-racing batch's other, still-running
# branches -- see _release_one_pending_for_call_id.
released_entry: Final = await self._release_one_pending_for_call_id(kwargs)
if released_entry is not None:
await self._release_keys((released_entry,))
async def async_log_success_event(
self,
@ -696,14 +779,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
start_time: datetime | None,
end_time: datetime | None,
) -> None:
stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs))
if stash is not None and stash.pending_concurrency_keys:
release_keys: Final = tuple(stash.pending_concurrency_keys)
stash.pending_concurrency_keys.clear()
release_task: Final = asyncio.create_task(self._release_keys(release_keys))
released_entry: Final = await self._release_one_pending_for_call_id(kwargs)
if released_entry is not None:
release_task: Final = asyncio.create_task(self._release_keys((released_entry,)))
_BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring
release_task.add_done_callback(_BACKGROUND_TASKS.discard)
stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs))
config: Final = self._refresh_config()
if config is None:
return

337
tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py
View file

@ -18,6 +18,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro
from litellm.proxy.hooks.global_tag_rate_limits_hook import (
_PROXY_GlobalTagRateLimitsHook,
)
from litellm.router import Router
class TimeController:
@ -60,7 +61,9 @@ def _redis_hook(time_controller: TimeController):
pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set")
redis_cache = RedisCache(host=redis_host, port=int(redis_port), password=os.getenv("REDIS_PASSWORD"))
dual_cache = DualCache(redis_cache=redis_cache)
return _PROXY_GlobalTagRateLimitsHook(internal_usage_cache=dual_cache, time_provider=time_controller.now), redis_cache
return _PROXY_GlobalTagRateLimitsHook(
internal_usage_cache=dual_cache, time_provider=time_controller.now
), redis_cache
# ---------------------------------------------------------------------------
@ -472,7 +475,9 @@ async def test_dollar_limit_respects_apply_to_models_at_accounting_time(time_con
@pytest.mark.asyncio
async def test_apply_to_models_fallback_does_not_re_narrow_accounting_to_the_serving_model(time_controller, monkeypatch):
async def test_apply_to_models_fallback_does_not_re_narrow_accounting_to_the_serving_model(
time_controller, monkeypatch
):
"""
Documented limitation, not a bug: apply_to_models is evaluated exactly
once, at admission, against the caller-requested model -- it is never
@ -673,7 +678,12 @@ async def test_a_rejected_admission_attempts_model_does_not_drive_later_accounti
"litellm_call_id": "call-1",
"metadata": {"tags": ["end_user_id:u1"]},
"model": "model-b",
"standard_logging_object": {"total_tokens": 0, "response_cost": 50.0, "model": "model-b", "model_group": "model-b"},
"standard_logging_object": {
"total_tokens": 0,
"response_cost": 50.0,
"model": "model-b",
"model_group": "model-b",
},
}
await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
@ -927,7 +937,11 @@ async def test_concurrency_reservation_released_when_a_different_hook_rejects_th
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{"concurrency_limits": {"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]}},
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
@ -1097,6 +1111,303 @@ async def test_concurrent_requests_do_not_share_each_others_reservation_state(ti
assert results == [1, 1]
# ---------------------------------------------------------------------------
# Non-racing batch dispatch (Router.abatch_completion via a comma-separated
# `model`): one admission fans out into several real, independent LLM calls,
# each firing its own terminal event -- confirmed live against the real
# Router.abatch_completion (not just this hook in isolation) that every
# branch always emits exactly one success or failure event, unlike the
# racing abatch_completion_fastest_response variant, whose cancelled losers
# never emit any terminal event at all.
# ---------------------------------------------------------------------------
def _batch_data(tags: list[str], model: str, call_id: str = "call-1", fastest_response: bool = False) -> dict:
data = {**_data(tags, call_id=call_id), "model": model}
if fastest_response:
data["fastest_response"] = True
return data
@pytest.mark.asyncio
async def test_non_racing_batch_reserves_one_unit_per_model(time_controller, monkeypatch):
"""A single admission for a 3-model comma-separated batch must reserve 3
units, not 1 -- otherwise the cap is measured against the wrong number of
real concurrent LLM calls the batch actually makes."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="acompletion",
)
# The batch alone already occupies every unit of the cap of 3.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-2"),
call_type="completion",
)
@pytest.mark.asyncio
async def test_non_racing_batch_releases_one_unit_per_branch(time_controller, monkeypatch):
"""Each branch's own success/failure event must release only its own
share of a non-racing batch's reservation -- releasing the whole
reservation on the first branch to finish would let a new, unrelated
request in while the batch's own remaining branches are still genuinely
in flight, exceeding the configured cap."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="acompletion",
)
batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}}
await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# Only 1 of 3 branches finished -- the other 2 still hold their own
# units, so a brand-new request needing all 3 remaining units of the cap
# must still be rejected.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-x,model-y,model-z", call_id="call-2"),
call_type="acompletion",
)
# The other 2 branches finish (one success, one failure).
await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await hook.async_log_failure_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# Now every unit from the first batch is back -- no leak, no
# over-release along the way.
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-3"),
call_type="completion",
)
assert result is not None
@pytest.mark.asyncio
async def test_non_racing_batch_disconnect_releases_only_what_remains(time_controller, monkeypatch):
"""A client disconnect mid-batch cancels every branch that hasn't
already finished together, so none of those cancelled branches ever
fires its own terminal event. The disconnect hook must release exactly
what's left -- not double-release a branch that already released its
own share via a real success/failure event first."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="acompletion",
)
batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}}
await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"})
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-2"),
call_type="completion",
)
assert result is not None
@pytest.mark.asyncio
async def test_fastest_response_batch_still_reserves_only_one_unit(time_controller, monkeypatch):
"""abatch_completion_fastest_response cancels every losing branch
without ever firing a terminal event for it, so it must keep reserving a
single unit for the whole dispatch -- reserving one per model here would
leak every losing branch's share until the safety TTL on every single
call."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b", call_id="call-1", fastest_response=True),
call_type="acompletion",
)
# If this admission had reserved 2 units instead of 1, this second,
# unrelated request (cap of 1) would already be rejected here too.
kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}}
await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-2"),
call_type="completion",
)
assert result is not None
@pytest.mark.asyncio
async def test_real_abatch_completion_admits_once_and_releases_per_branch(time_controller, monkeypatch):
"""End-to-end against the real Router.abatch_completion (mocked LLM
responses, no network): confirms admission fires exactly once for the
whole dispatch -- not once per branch the way model_based_tag_rate_limits_hook's
per-Router-hop admission does -- and that a cap sized to the batch width
is never exceeded even while some branches are still genuinely running."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
monkeypatch.setattr(litellm, "callbacks", [hook])
router = Router(
model_list=[
{
"model_name": "model-fast",
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "fast-done", "mock_delay": 0.01},
},
{
"model_name": "model-slow-1",
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-1-done", "mock_delay": 0.3},
},
{
"model_name": "model-slow-2",
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-2-done", "mock_delay": 0.3},
},
]
)
admitted_data = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-fast,model-slow-1,model-slow-2", call_id="call-1"),
call_type="acompletion",
)
batch_task = asyncio.create_task(
router.abatch_completion(
models=["model-fast", "model-slow-1", "model-slow-2"],
messages=[{"role": "user", "content": "hello"}],
metadata=admitted_data["metadata"],
litellm_call_id=admitted_data["litellm_call_id"],
)
)
# Give only the fast branch (0.01s mock_delay) time to finish; the two
# slow branches (0.3s) are still genuinely in flight.
await asyncio.sleep(0.1)
# 1 of 3 units released, 2 still held by the still-running slow
# branches -- a brand-new request needing all 3 units of the cap must
# still be rejected.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-x,model-y,model-z", call_id="call-2"),
call_type="acompletion",
)
responses = await batch_task
assert [r.choices[0].message.content for r in responses] == ["fast-done", "slow-1-done", "slow-2-done"]
await asyncio.sleep(0)
# Every branch has now finished and released its own share -- no leak.
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-3"),
call_type="completion",
)
assert result is not None
@pytest.mark.asyncio
async def test_comma_separated_model_outside_acompletion_call_type_is_not_a_batch(time_controller, monkeypatch):
"""A comma in `model` only means a non-racing batch dispatch for the
exact call_type route_llm_request.py itself gates on -- any other
call_type must reserve just 1 unit, matching whatever that call_type
actually does under the hood."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="completion",
)
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-2"),
call_type="completion",
)
# ---------------------------------------------------------------------------
# Accounting: tokens/dollars via async_log_success_event
# ---------------------------------------------------------------------------
@ -1368,7 +1679,9 @@ async def test_a_caller_supplied_tag_cannot_shadow_the_policy_backed_identity_ta
"metadata": {"tags": ["company_id:real-company"], "inherited_tags": ["company_id:real-company"]},
}
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=victim_data, call_type="completion")
await hook.async_pre_call_hook(
user_api_key_dict=key, cache=DualCache(), data=victim_data, call_type="completion"
)
@pytest.mark.asyncio
@ -1380,7 +1693,9 @@ async def test_success_accounting_also_resolves_identity_from_the_policy_backed_
"global_tag_rate_limits",
{
"dollar_limits": {
"limits": [{"name": "per-company-spend", "tag_id": "company_id", "limit": 10.0, "period_seconds": 86400}]
"limits": [
{"name": "per-company-spend", "tag_id": "company_id", "limit": 10.0, "period_seconds": 86400}
]
}
},
)
@ -1440,7 +1755,10 @@ async def test_rejection_detail_does_not_disclose_the_resolved_tag_value(time_co
data = {
"litellm_call_id": "call-1",
"metadata": {"tags": ["company_id:secret-internal-name"], "inherited_tags": ["company_id:secret-internal-name"]},
"metadata": {
"tags": ["company_id:secret-internal-name"],
"inherited_tags": ["company_id:secret-internal-name"],
},
}
await hook.async_pre_call_hook(user_api_key_dict=key, cache=DualCache(), data=data, call_type="completion")
@ -1583,7 +1901,10 @@ async def test_admission_ignores_a_forged_empty_litellm_metadata_key(time_contro
poisoned = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": {}}
await hook.async_pre_call_hook(
user_api_key_dict=_key(), cache=DualCache(), data={**poisoned, "litellm_call_id": "call-1"}, call_type="completion"
user_api_key_dict=_key(),
cache=DualCache(),
data={**poisoned, "litellm_call_id": "call-1"},
call_type="completion",
)
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(