From 4d5f1dd144d6653ba3f16aa4a58c4e50e45a8fc0 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 6 Nov 2026 04:06:57 +0000 Subject: [PATCH] fix(proxy): charge actual tokens to the per-customer key counter after the call, report customer limit edits as dirty fields --- .../hooks/parallel_request_limiter_v3.py | 51 +++++++++++------- .../hooks/test_parallel_request_limiter_v3.py | 54 ++++++++++--------- .../CustomerRateLimitEditor.test.tsx | 21 ++++++++ .../CustomerRateLimitEditor.tsx | 7 +++ .../key_edit_view.integration.test.tsx | 6 +++ .../components/templates/key_edit_view.tsx | 19 +++++-- 6 files changed, 110 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b964d789adb..d64a2a7e2f9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -761,6 +761,7 @@ class RequestRateLimiterStash: reservation_released: bool = False io_reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) + customer_per_key_tpm_scope: tuple[str, str] | None = None socket_agent_windows: tuple[RateLimitDescriptor, ...] = () def clear_reserved_tokens(self) -> None: @@ -3219,13 +3220,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - def _add_customer_per_key_rate_limit_descriptor( + def _customer_per_key_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, - descriptors: list[RateLimitDescriptor], - ) -> None: + ) -> RateLimitDescriptor | None: """ - Add a per-customer rpm/tpm descriptor scoped to the API key. + Per-customer rpm/tpm descriptor scoped to the API key. The counter is ``{api_key}:{customer}``, so one customer exhausting its allowance on this key never blocks the other customers on the key, nor @@ -3234,23 +3234,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ customer: Final = user_api_key_dict.end_user_id if not user_api_key_dict.api_key or not customer: - return + return None rpm_limit: Final = get_key_customer_rpm_limit(user_api_key_dict).get(customer) tpm_limit: Final = get_key_customer_tpm_limit(user_api_key_dict).get(customer) if rpm_limit is None and tpm_limit is None: - return + return None - descriptors.append( - RateLimitDescriptor( - key="customer_per_key", - value=f"{user_api_key_dict.api_key}:{customer}", - rate_limit={ - "requests_per_unit": rpm_limit, - "tokens_per_unit": tpm_limit, - "window_size": self.window_size, - }, - ) + return RateLimitDescriptor( + key="customer_per_key", + value=f"{user_api_key_dict.api_key}:{customer}", + rate_limit={ + "requests_per_unit": rpm_limit, + "tokens_per_unit": tpm_limit, + "window_size": self.window_size, + }, ) def _add_mcp_per_key_rate_limit_descriptor( @@ -3671,10 +3669,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors=descriptors, ) - self._add_customer_per_key_rate_limit_descriptor( - user_api_key_dict=user_api_key_dict, - descriptors=descriptors, - ) + customer_descriptor: Final = self._customer_per_key_rate_limit_descriptor(user_api_key_dict=user_api_key_dict) + if customer_descriptor is not None: + descriptors.append(customer_descriptor) self._add_team_model_rate_limit_descriptor_from_metadata( user_api_key_dict=user_api_key_dict, @@ -4500,6 +4497,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for d in descriptors if d["key"] == "tag" and d["rate_limit"] is not None and d["rate_limit"].get("tokens_per_unit") is not None ) + stash.customer_per_key_tpm_scope = next( + ( + (d["key"], d["value"]) + for d in descriptors + if d["key"] == "customer_per_key" + and d["rate_limit"] is not None + and d["rate_limit"].get("tokens_per_unit") is not None + ), + None, + ) stash.socket_agent_windows = tuple(d for d in descriptors if d["key"] in _AGENT_DESCRIPTOR_KEYS) # Only check rate limits if we have descriptors with actual limits @@ -5282,6 +5289,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): kwargs: object, model_group: str | None, tpm_limited_tags: Set[str] = frozenset(), + customer_per_key_scope: tuple[str, str] | None = None, ) -> list[tuple[str, str]]: """ Enumerate every (scope_key, scope_value) pair that *might* carry a @@ -5338,6 +5346,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if session_id: targets.append(("agent_session", f"{agent_id}:{session_id}")) targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags)) + if customer_per_key_scope is not None: + targets.append(customer_per_key_scope) return targets def _build_reservation_aware_tpm_ops( @@ -5514,6 +5524,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): standard_logging_metadata=standard_logging_metadata, kwargs=kwargs, tpm_limited_tags=_tpm_limited_tags(stash, kwargs), + customer_per_key_scope=stash.customer_per_key_tpm_scope if stash is not None else None, model_group=reconcile_model.group if reconcile_model is not None else None, ) targets.extend( diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 7fc39d3cb1f..325746fbe03 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -10131,33 +10131,39 @@ async def test_per_customer_rpm_limit_on_key_is_scoped_to_that_customer_v3(monke await call(other_key, "acme") -def test_per_customer_descriptor_carries_rpm_and_tpm_for_the_requesting_customer_v3(): - """ - create_rate_limit_descriptors emits one customer_per_key descriptor keyed by - ``{api_key}:{customer}`` with that customer's rpm and tpm, and none when the - requesting customer has no configured limit. - """ - api_key = hash_token("sk-per-customer-desc") - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - metadata = {"customer_rpm_limit": {"acme": 5}, "customer_tpm_limit": {"acme": 900, "globex": 50}} +@pytest.mark.asyncio +async def test_per_customer_tpm_limit_on_key_charges_actual_usage_after_the_call_v3(monkeypatch): + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") + _request_stash.set(None) + api_key = hash_token("sk-per-customer-tpm") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") + metadata = {"customer_tpm_limit": {"acme": 100}} - def descriptors_for(customer: str) -> list[RateLimitDescriptor]: - descriptors = handler.create_rate_limit_descriptors( + async def call(customer: str) -> None: + await handler.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(api_key=api_key, end_user_id=customer, metadata=metadata), + cache=local_cache, data={"model": "gpt-3.5-turbo", "user": customer}, - rpm_limit_type=None, - tpm_limit_type=None, - model_has_failures=False, + call_type="", ) - return [d for d in descriptors if d["key"] == "customer_per_key"] - (acme,) = descriptors_for("acme") - assert acme["value"] == f"{api_key}:acme" - assert acme["rate_limit"]["requests_per_unit"] == 5 - assert acme["rate_limit"]["tokens_per_unit"] == 900 + await call("acme") + await handler.async_log_success_event( + kwargs={ + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + "model": "gpt-3.5-turbo", + "user": "acme", + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)), + start_time=datetime.now(), + end_time=datetime.now(), + ) - (globex,) = descriptors_for("globex") - assert globex["rate_limit"]["requests_per_unit"] is None - assert globex["rate_limit"]["tokens_per_unit"] == 50 - - assert descriptors_for("initech") == [] + with pytest.raises(HTTPException) as exc_info: + await call("acme") + assert exc_info.value.status_code == 429 + assert "customer_per_key" in str(exc_info.value.detail) + await call("globex") diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.test.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.test.tsx index 5d09a573001..0292528b5f6 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.test.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.test.tsx @@ -7,6 +7,7 @@ import { CustomerRateLimitEntry, customerLimitsToRows, customerRowsToLimits, + customerLimitsEqual, } from "./CustomerRateLimitEditor"; function Harness({ @@ -101,3 +102,23 @@ describe("customerLimitsToRows", () => { expect(customerLimitsToRows(undefined, null)).toStrictEqual([]); }); }); + +describe("customerLimitsEqual", () => { + it("treats the same limits in a different row order as unchanged", () => { + expect( + customerLimitsEqual( + { customer_rpm_limit: { acme: 2, globex: 5 }, customer_tpm_limit: {} }, + { customer_rpm_limit: { globex: 5, acme: 2 }, customer_tpm_limit: {} }, + ), + ).toBe(true); + }); + + it("reports a changed, added or removed customer limit", () => { + const stored = { customer_rpm_limit: { acme: 2 }, customer_tpm_limit: { acme: 900 } }; + expect(customerLimitsEqual(stored, { customer_rpm_limit: { acme: 3 }, customer_tpm_limit: { acme: 900 } })).toBe( + false, + ); + expect(customerLimitsEqual(stored, { customer_rpm_limit: { acme: 2 }, customer_tpm_limit: {} })).toBe(false); + expect(customerLimitsEqual(stored, { ...stored, customer_rpm_limit: { acme: 2, globex: 1 } })).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.tsx index 6af43913915..ff026fdf024 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.tsx @@ -39,6 +39,13 @@ const toNumberMap = (raw: unknown): Record => { return out; }; +const sameLimitMap = (a: Record, b: Record): boolean => + Object.keys(a).length === Object.keys(b).length && + Object.entries(a).every(([customer, limit]) => b[customer] === limit); + +export const customerLimitsEqual = (a: CustomerRateLimits, b: CustomerRateLimits): boolean => + sameLimitMap(a.customer_rpm_limit, b.customer_rpm_limit) && sameLimitMap(a.customer_tpm_limit, b.customer_tpm_limit); + export const customerLimitsToRows = ( customerRpmLimit?: unknown, customerTpmLimit?: unknown, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index 668ff955f4d..a3d968f4e53 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -2592,6 +2592,8 @@ describe("KeyEditView", () => { auto_rotate: false, rotation_interval: undefined, tag_rpm_limit: {}, + customer_rpm_limit: {}, + customer_tpm_limit: {}, }; describe("submit payload contract", () => { @@ -3171,6 +3173,7 @@ describe("KeyEditView", () => { }); expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("customer_rpm_limit", { acme: 7 }); expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("customer_tpm_limit", { acme: 900 }); + expect(onSubmitMock.mock.calls[0][1]).not.toContain("customer_rpm_limit"); }); it("clears a stored customer limit when its row is removed", async () => { @@ -3188,6 +3191,9 @@ describe("KeyEditView", () => { expect(onSubmitMock).toHaveBeenCalled(); }); expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("customer_rpm_limit", {}); + expect(onSubmitMock.mock.calls[0][1]).toEqual( + expect.arrayContaining(["customer_rpm_limit", "customer_tpm_limit"]), + ); }); const setRpmLimit = (value: string) => { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index c18f43c1edb..38cfbd6ce92 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -68,6 +68,7 @@ import { import { CustomerRateLimitEditor, CustomerRateLimitEntry, + customerLimitsEqual, customerLimitsToRows, customerRowsToLimits, } from "../key_team_helpers/CustomerRateLimitEditor"; @@ -303,9 +304,15 @@ function KeyEditSession({ // clears the stored limits ({} overwrites the metadata field). const { tag_rpm_limit } = tagRowsToLimits(tagRateLimits); values.tag_rpm_limit = tag_rpm_limit; - const { customer_rpm_limit, customer_tpm_limit } = customerRowsToLimits(customerRateLimits); - values.customer_rpm_limit = customer_rpm_limit; - values.customer_tpm_limit = customer_tpm_limit; + const customerLimits = customerRowsToLimits(customerRateLimits); + values.customer_rpm_limit = customerLimits.customer_rpm_limit; + values.customer_tpm_limit = customerLimits.customer_tpm_limit; + const customerLimitsUnchanged = customerLimitsEqual( + customerLimits, + customerRowsToLimits( + customerLimitsToRows(keyData.metadata?.customer_rpm_limit, keyData.metadata?.customer_tpm_limit), + ), + ); const hadExistingFallbacks = keyData.budget_fallbacks != null && Object.keys(keyData.budget_fallbacks).length > 0; if (Object.keys(budgetFallbacks).length > 0) { @@ -334,7 +341,11 @@ function KeyEditSession({ ...values, ...(detachProject && enableProjectsUI && canDetachProject ? { project_id: null } : {}), }), - [...Object.keys(dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])], + [ + ...Object.keys(dirtyFields), + ...(budgetLimitsUnchanged ? [] : ["budget_limits"]), + ...(customerLimitsUnchanged ? [] : ["customer_rpm_limit", "customer_tpm_limit"]), + ], ); } finally { setIsKeySaving(false);