fix(proxy): size the uncapped output floor against the PTU share in raw output tokens and escape the exported team id

This commit is contained in:
mateo-berri 2026-09-24 16:06:01 -07:00
parent 0d80930df8
commit 4ad0db30b2
5 changed files with 71 additions and 4 deletions

View file

@ -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

View file

@ -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 "-",

View file

@ -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]):

View file

@ -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()

View file

@ -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