From a3d791f34858abd06bba96e09d9c9330b86a9240 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 24 Sep 2026 02:49:01 +0000
Subject: [PATCH] test: add tag rpm limit and tag budget reset integration
 coverage (#42859)

Adds two integration tests to tests/integration/spend/test_tag_budget_enforcement.py

test_key_tag_rpm_limit_rejects_the_second_request_carrying_that_tag proves
a key metadata tag_rpm_limit of 1 rejects the second request carrying that
tag with 429 while a request carrying a different tag still passes

test_tag_budget_duration_resets_spend_and_unblocks_the_tag boots an owned
proxy with a 2 to 3 second budget rescheduler, creates a tag with
max_budget 0.0001 and budget_duration 5s, observes the spend block, then
observes the tag serving again once ResetBudgetJob zeroes the tag spend

Mutation evidence

get_key_tag_rpm_limit forced to return None: the second tagged request
returned 200 instead of 429, test red

_queue_budget_linked_resets for uow.tags disabled in
_commit_budget_cascade_once: the tag stayed blocked at 422 for the full
70 second recovery window, test red

Co-authored-by: kerry <kerry@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../spend/test_tag_budget_enforcement.py      | 98 +++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/tests/integration/spend/test_tag_budget_enforcement.py b/tests/integration/spend/test_tag_budget_enforcement.py
index e8d4c6438a5..f29ee6c07c8 100644
--- a/tests/integration/spend/test_tag_budget_enforcement.py
+++ b/tests/integration/spend/test_tag_budget_enforcement.py
@@ -1,7 +1,9 @@
 import uuid
+from pathlib import Path
 from typing import Final
 
 from integration._support.client import Gateway, eventually
+from integration._support.process import owned_proxy
 
 
 def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway) -> None:
@@ -58,3 +60,99 @@ def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway)
             },
         )
         assert control.status_code == 200, control.text
+
+
+def test_key_tag_rpm_limit_rejects_the_second_request_carrying_that_tag(gateway: Gateway) -> None:
+    tag: Final = f"tag-rpm-{uuid.uuid4().hex}"
+    with gateway.scenario() as scenario:
+        model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
+        key: Final = scenario.key(metadata={"tag_rpm_limit": {tag: 1}})
+        first: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": f"tag rpm {tag}"}],
+                "metadata": {"tags": [tag]},
+            },
+            key=key,
+        )
+        assert first.status_code == 200, first.text
+        second: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": f"tag rpm {tag}"}],
+                "metadata": {"tags": [tag]},
+            },
+            key=key,
+        )
+        assert second.status_code == 429, second.text
+        assert "rpm" in second.text.lower() or "rate" in second.text.lower(), second.text
+        control: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": f"other tag rpm {tag}"}],
+                "metadata": {"tags": [f"other-{tag}"]},
+            },
+            key=key,
+        )
+        assert control.status_code == 200, control.text
+
+
+def test_tag_budget_duration_resets_spend_and_unblocks_the_tag(gateway: Gateway, tmp_path: Path) -> None:
+    tag: Final = f"tag-reset-{uuid.uuid4().hex}"
+    with (
+        owned_proxy(
+            gateway,
+            tmp_path,
+            {"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "3"},
+        ) as candidate,
+        candidate.scenario() as scenario,
+    ):
+
+        def delete_tag() -> None:
+            candidate.post("/tag/delete", {"name": tag})
+
+        model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01)
+        candidate.post("/tag/new", {"name": tag, "max_budget": 0.0001, "budget_duration": "5s"})
+        scenario.cleanups.callback(delete_tag)
+        first: Final = candidate.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": f"tag spend {tag}"}],
+                "metadata": {"tags": [tag]},
+            },
+        )
+        assert first.status_code == 200, first.text
+
+        def rejection() -> int:
+            return candidate.request(
+                "POST",
+                "/v1/chat/completions",
+                {
+                    "model": model,
+                    "messages": [{"role": "user", "content": f"tag reset probe {tag}"}],
+                    "metadata": {"tags": [tag]},
+                },
+            ).status_code
+
+        status: Final = eventually(rejection, lambda code: code != 200, seconds=70)
+        assert status in (400, 422, 429), status
+        blocked: Final = candidate.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": f"tag reset probe {tag}"}],
+                "metadata": {"tags": [tag]},
+            },
+        )
+        assert "budget" in blocked.text.lower(), blocked.text
+        recovered: Final = eventually(rejection, lambda code: code == 200, seconds=70)
+        assert recovered == 200, recovered