From 4ad0db30b26c94a7d9295a99858ed8e3f62e53e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:06:01 -0700 Subject: [PATCH] fix(proxy): size the uncapped output floor against the PTU share in raw output tokens and escape the exported team id --- .../hooks/parallel_request_limiter_v3.py | 11 +++++-- .../management_endpoints/team_endpoints.py | 2 +- litellm/router_utils/ptu_shares.py | 4 +++ .../hooks/test_parallel_request_limiter_v3.py | 26 +++++++++++++++ .../test_team_endpoints.py | 32 +++++++++++++++++++ 5 files changed, 71 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c61b35f5936..e0f0b9d714e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -3765,13 +3765,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # in-memory check otherwise — single-worker protection still holds # even without Redis. # ---------------------------------------------------------------- - configured_tpm_limits: Final = [ - int(v) + ptu_raw_output_limit: Final = ( + stash.ptu_ceiling.raw_output_token_limit() if stash.ptu_ceiling is not None else None + ) + configured_tpm_limits: Final = tuple( + ptu_raw_output_limit + if d["key"] == PTU_TEAM_DESCRIPTOR_KEY and ptu_raw_output_limit is not None + else int(v) for d in descriptors if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] if v is not None - ] + ) has_tpm_limits: Final = bool(configured_tpm_limits) # Populated on a successful combined-TPM reservation below, so the diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b3dca8b127f..b0320cd6e6d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -6882,7 +6882,7 @@ def _export_csv_record(row: TeamDailyActivityExportRow) -> dict[str, object]: return { # mutable-ok: csv.DictWriter consumes a plain mapping per row "Date": row.date, "Team": _csv_safe(row.team_alias) if row.team_alias else "-", - "Team ID": row.team_id, + "Team ID": _csv_safe(row.team_id), "Key Alias": _csv_safe(row.key_alias) if row.key_alias else "-", "Key ID": row.api_key or "-", "User ID": _csv_safe(row.user_id) if row.user_id else "-", diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index a797cd2d109..3bc9d33349c 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -25,6 +25,10 @@ class PTUTeamCeiling: output_to_input_ratio: float cached_input_ratio: float + def raw_output_token_limit(self) -> int: + """The ceiling expressed in unweighted tokens: what fits under it when every token is output.""" + return max(1, int(self.tpm_limit / max(self.output_to_input_ratio, 1.0))) + @dataclass(frozen=True, slots=True) class PTUShareFilterResult(Generic[_DeploymentT]): diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 09a8314dc37..dee97db8472 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7285,6 +7285,32 @@ async def test_the_reservation_weighs_output_the_way_the_ceiling_does(): assert "model_per_team_ptu" in str(rejected.value.detail) +@pytest.mark.asyncio +async def test_a_one_ptu_share_admits_four_uncapped_requests_a_minute_and_rejects_the_fifth(): + """One PTU of gpt-4.1 is 3,000 normalized tokens a minute. A request without max_tokens gets the + proxy's output floor measured against the share in output tokens (750), so its cap costs at most + a quarter of the share once weighted 4:1 and four such requests fit before the fifth is refused.""" + cache = DualCache() + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=3000, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + uncapped = {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]} + + admitted = [dict(uncapped) for _ in range(4)] + for data in admitted: + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=data, call_type="acompletion") + with pytest.raises(HTTPException) as fifth: + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=dict(uncapped), call_type="acompletion" + ) + + assert fifth.value.status_code == 429 + assert "model_per_team_ptu" in str(fifth.value.detail) + assert all(data["max_tokens"] * 4 <= 3000 // 4 for data in admitted) + + @pytest.mark.asyncio async def test_the_ptu_counter_holds_the_normalized_reservation_beside_the_raw_one(): cache = DualCache() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 87fff1d5d6a..6d87eda3d98 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -17295,6 +17295,38 @@ def test_team_export_csv_columns_match_the_dashboard_client_layout(): ] +def test_team_export_csv_escapes_a_formula_shaped_team_id_like_the_alias(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="=HYPERLINK(\"https://example.com\")", + team_alias="=SUM(A1)", + api_key=None, + key_alias=None, + user_id=None, + user_email=None, + spend=0.0, + api_requests=0, + successful_requests=0, + failed_requests=0, + total_tokens=0, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ) + + record: Final = next(csv.DictReader(io.StringIO(_team_export_csv("daily", (row,))))) + + assert record["Team ID"] == "'=HYPERLINK(\"https://example.com\")" + assert record["Team"] == "'=SUM(A1)" + + def test_team_export_csv_omits_key_columns_for_the_plain_daily_scope(): import csv import io