fix(proxy): avoid duplicate model_per_team descriptors in v3 limiter

Team model TPM/RPM from metadata was merged into descriptors twice
(_create_rate_limit_descriptors plus async_pre_call_hook), doubling
increments on the same model_per_team counter. Remove redundant path
and add regression tests.

Made-with: Cursor
This commit is contained in:
Milan 2026-04-30 21:25:02 +03:00
parent 609454d0f1
commit 6909fd248a
No known key found for this signature in database
2 changed files with 63 additions and 60 deletions

View file

@ -1122,59 +1122,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return batch_limiter
return None
def _add_team_model_rate_limit_descriptor_from_metadata(
self,
user_api_key_dict: UserAPIKeyAuth,
requested_model: Optional[str],
descriptors: List[RateLimitDescriptor],
) -> None:
"""Add team model rate limit descriptor from team_metadata if applicable."""
if (
get_model_rate_limit_from_metadata(
user_api_key_dict, "team_metadata", "model_rpm_limit"
)
is not None
or get_model_rate_limit_from_metadata(
user_api_key_dict, "team_metadata", "model_tpm_limit"
)
is not None
):
_tpm_limit_for_team_model = (
get_model_rate_limit_from_metadata(
user_api_key_dict, "team_metadata", "model_tpm_limit"
)
or {}
)
_rpm_limit_for_team_model = (
get_model_rate_limit_from_metadata(
user_api_key_dict, "team_metadata", "model_rpm_limit"
)
or {}
)
should_check_rate_limit = (
requested_model in _tpm_limit_for_team_model
or requested_model in _rpm_limit_for_team_model
)
if should_check_rate_limit and requested_model is not None:
model_specific_tpm_limit = _tpm_limit_for_team_model.get(
requested_model
)
model_specific_rpm_limit = _rpm_limit_for_team_model.get(
requested_model
)
descriptors.append(
RateLimitDescriptor(
key="model_per_team",
value=f"{user_api_key_dict.team_id}:{requested_model}",
rate_limit={
"requests_per_unit": model_specific_rpm_limit,
"tokens_per_unit": model_specific_tpm_limit,
"window_size": self.window_size,
},
)
)
def _add_project_model_rate_limit_descriptor_from_metadata(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -1332,13 +1279,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
model_has_failures=model_has_failures,
)
# Add team model rate limits from team_metadata
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
# Project Level Rate Limits
self._add_project_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,

View file

@ -2685,3 +2685,66 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3():
assert (
"model_per_project" not in descriptor_keys
), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}"
def test_team_metadata_model_tpm_limit_single_model_per_team_descriptor():
"""
team_metadata.model_tpm_limit must not produce duplicate model_per_team descriptors.
Duplicates caused the same counter to increment twice per request (effective limit halved).
"""
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user = UserAPIKeyAuth(
api_key=hash_token("sk-team-model-tpm-dedupe"),
team_id="team-dedupe-test",
team_metadata={"model_tpm_limit": {"gpt-4o-mini": 100}},
)
data = {"model": "gpt-4o-mini"}
descriptors = parallel_request_handler._create_rate_limit_descriptors(
user_api_key_dict=user,
data=data,
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
)
mpt = [d for d in descriptors if d["key"] == "model_per_team"]
assert len(mpt) == 1, (
f"expected exactly one model_per_team descriptor, got {len(mpt)}: {mpt}"
)
assert mpt[0]["value"] == "team-dedupe-test:gpt-4o-mini"
assert mpt[0]["rate_limit"]["tokens_per_unit"] == 100
@pytest.mark.asyncio
async def test_team_metadata_model_tpm_pre_call_not_double_incremented():
"""
With team model TPM = N, pre-call should allow N increments on model_per_team
(one per request), not N/2 from duplicate descriptors.
"""
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
user = UserAPIKeyAuth(
api_key=hash_token("sk-team-model-tpm-pre"),
team_id="team-pre-dedupe",
team_metadata={"model_tpm_limit": {"gpt-4o-mini": 10}},
)
data = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "x"}]}
allowed = 0
for _ in range(20):
try:
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user,
cache=local_cache,
data=dict(data),
call_type="completion",
)
allowed += 1
except HTTPException:
break
assert allowed == 10, (
f"expected 10 pre-calls before model_per_team TPM block, got {allowed}"
)