fix(proxy): don't release a live sibling's slot when this hop's own context resolves to nothing configured, and lock in multi-policy release

Cursor Bugbot: _release_own_concurrency_keys still fell back to an
unconditional release whenever _resolve_hop_context returned None, which
included the common case of a hop whose own model has no tag_rate_limits
at all. That hop's own admission never queues a reservation either (it
hits the identical "nothing configured" check first), so on
abatch_completion the unconditional release could still sweep up a
genuinely live sibling's own slot.

_resolve_hop_context now folds an empty configured/tags into a
resolved-but-empty context instead of returning None for those two cases
specifically, since _own_concurrency_keys_for_hop already computes an
empty release set from them on its own. None is reserved for the
genuinely ambiguous causes (no router, no standard_logging_object, no
model_group) where this hop's own admission could have reserved something
real but there's no data left to recompute it -- e.g. a different
registered CustomLogger rejecting the request before this hop's own call
ever ran. Those still need the pre-existing unconditional fallback, which
three existing tests already depended on and caught immediately when the
first version of this fix collapsed both cases together.

Also adds direct empirical coverage (prompted by the sibling PR's global
hook finding) confirming a hop matching two different concurrency-scoped
entries releases both reservations at once, not just one.
This commit is contained in:
Deepanshu 2026-09-03 10:38:58 -04:00
parent 180b5fd69b
commit f825c03b67
2 changed files with 169 additions and 21 deletions

View file

@ -1969,17 +1969,25 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
kwargs, model_group, fallback=live_candidate_model_names
)
configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names)
if not configured:
return None
tags: Final = _order_tags_for_identity_resolution(
_get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name),
kwargs,
metadata_variable_name,
)
if not tags:
return None
# Deliberately not an early `return None` when `configured` or
# `tags` is empty (Cursor Bugbot finding: that used to make
# `_release_own_concurrency_keys` fall back to an unconditional
# release): `async_filter_deployments` hits this exact same
# "nothing configured" / "no matching tag" check on *this hop's own
# admission* before it could ever queue a concurrency reservation,
# so an empty `configured` or `tags` here means this hop reserved
# nothing of its own to release either -- `_own_concurrency_keys_for_hop`
# already returns an empty set for either case on its own, without
# needing this function to special-case it. `None` is reserved for
# the genuinely ambiguous causes above (no router, no
# `standard_logging_object`, no `model_group`), where this hop's own
# admission *could* have reserved something we have no way to
# recompute -- those still need the unconditional fallback.
return _HopContext(
standard_logging_object=standard_logging_object,
configured=configured,
@ -1999,22 +2007,28 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
computed them -- never "every reservation currently pending", which
would also release a still-live `abatch_completion` sibling branch's
own reservation (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring).
Falls back to the pre-existing unconditional release only when
`context` itself couldn't be resolved (nothing configured for this
tag/model, or an identity-extraction edge case) -- in which case there
is no way to tell a sibling's reservation apart from this hop's own
anyway, so this only ever matters for the single-branch case that
release already handled correctly before `abatch_completion` existed.
A resolved `context` whose own hop reserved *no* concurrency slot at
all (no concurrency-unit limit matches its model/tags/deployment --
e.g. only a token or dollar limit is configured for this tag) must
release nothing, not fall through to the unconditional path: on
`abatch_completion` that shared pending list can still hold a
genuinely live sibling branch's own reservation (a different
deployment/model_group whose admission did reserve one), and this
hop reserving zero keys is not the same signal as `context` failing
to resolve at all.
A resolved `context` whose own hop matched no concurrency-unit limit
at all (Cursor Bugbot finding: "no `tag_rate_limits` configured for
this model" is the common case) must release nothing, not fall
through to the unconditional path below: `async_filter_deployments`
hits that identical "nothing configured"/"no matching tag" check on
*this exact hop's own admission*, before it could ever queue a
concurrency reservation, so this hop reserved nothing of its own to
release either -- `_resolve_hop_context` folds both into an
empty-but-resolved `context` rather than `None` for exactly this
reason (see its own docstring).
`context is None` means `_resolve_hop_context` hit one of its
genuinely ambiguous causes instead (no router, no
`standard_logging_object`, no `model_group`) -- cases where this
hop's own admission *could* have reserved something, but there is
no reliable data left to recompute what. Only those fall back to
the pre-existing unconditional release, the same as before
`abatch_completion` existed: e.g. a *different* registered
CustomLogger rejecting the request before this hop's own call ever
ran (so no `standard_logging_object` was ever built) must still
release this hook's own successfully reserved slot.
"""
if context is None:
return await self._pop_pending_concurrency_keys(kwargs)

View file

@ -3389,6 +3389,140 @@ async def test_batch_sibling_reserving_no_concurrency_slot_does_not_release_a_st
)
@pytest.mark.asyncio
async def test_batch_sibling_with_no_tag_rate_limits_at_all_does_not_release_a_still_executing_siblings_slot(
time_controller,
):
"""
Cursor Bugbot finding: `_release_own_concurrency_keys` still fell back to
an unconditional release whenever `_resolve_hop_context` returned `None`
-- including the common case of a hop whose own model has no
`tag_rate_limits` configured at all, so `resolve_any` finds nothing and
`_resolve_hop_context` bails out. That same "nothing configured" check
is also the *first* thing this hop's own admission does, before it could
ever queue a concurrency reservation -- so a `None` context means this
hop reserved nothing of its own, the same as a resolved context with no
matching concurrency-unit limit. On `abatch_completion` the still-live
sibling's own reservation must survive this branch's completion.
"""
limiter = _make_limiter(time_controller)
router = litellm.Router(
model_list=[
_deployment(
"grp-conc",
"dep-conc",
{
"concurrency_limits": {
"limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}]
}
},
),
_deployment("grp-none", "dep-none", {}),
]
)
limiter.update_variables(llm_router=router)
request_kwargs, kwargs = _call_context(["end_user_id:u1"])
conc_healthy = [d for d in router.model_list if d["model_info"]["id"] == "dep-conc"]
none_healthy = [d for d in router.model_list if d["model_info"]["id"] == "dep-none"]
async def _admit(model: str, healthy: list) -> None:
await limiter.async_filter_deployments(
model=model, healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
)
# Both branches of one abatch_completion dispatch admit under the
# identical shared model_call_details, each its own asyncio.Task -- see
# the concurrency-sibling tests above for why that models two
# independent admission lineages.
await asyncio.create_task(_admit("grp-conc", conc_healthy))
await asyncio.create_task(_admit("grp-none", none_healthy))
# The unconfigured branch finishes first. Its own success event resolves
# no context at all (nothing configured for "grp-none"), and reserved no
# concurrency slot of its own -- it must not touch the sibling's.
kwargs["standard_logging_object"] = {
"model_group": "grp-none",
"model_id": "dep-none",
"total_tokens": 10,
"response_cost": 0.01,
}
await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# The concurrency-limited branch's own slot must still be held: a fresh
# request against the same tag/model is rejected under the limit=1 cap.
with pytest.raises(ProxyRateLimitError):
await limiter.async_filter_deployments(
model="grp-conc",
healthy_deployments=conc_healthy,
messages=None,
request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}},
)
@pytest.mark.asyncio
async def test_a_hop_matching_two_concurrency_scoped_entries_releases_both_reservations(time_controller):
"""
#38347's global hook had a real bug on completely ordinary, non-batch
traffic: a request matching multiple concurrency-scoped TagRateLimitEntry
policies reserved one key per matching policy at admission but a
terminal event released only one entry total, leaking every other match
until its safety TTL. Verifies this hook's own release path -- which
computes `only_keys` as the full set of every matching policy's own key,
not a single arbitrary one -- actually releases all of them, not just
the first.
"""
limiter = _make_limiter(time_controller)
router = litellm.Router(
model_list=[
_deployment(
"grp",
"dep-1",
{
"concurrency_limits": {
"limits": [
{"name": "per-user", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300},
{"name": "per-team", "tag_id": "team_id", "limit": 1, "period_seconds": 300},
]
}
},
)
]
)
limiter.update_variables(llm_router=router)
request_kwargs, kwargs = _call_context(["end_user_id:u1", "team_id:t1"])
healthy = router.model_list
admitted = await limiter.async_filter_deployments(
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
)
assert admitted == healthy
pending = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD)
assert pending is not None and len(pending) == 2
kwargs["standard_logging_object"] = {
"model_group": "grp",
"model_id": "dep-1",
"total_tokens": 10,
"response_cost": 0.01,
}
await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
assert not kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD)
# Isolated per-policy checks: a follow-up request carrying only one tag
# exercises only that one policy, so each must independently show its
# own limit=1 slot free again, not just "no exception" from whichever
# one entry a partial release happened to free.
await limiter.async_filter_deployments(
model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}
)
await limiter.async_filter_deployments(
model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["team_id:t1"]}}
)
def _request_limit_router(limit: int) -> "litellm.Router":
return litellm.Router(
model_list=[