From 9b42809a816b5da682a2c9b296baf318231c4202 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 22 Sep 2026 21:26:48 +0000 Subject: [PATCH] test(integration): fail-closed key rejects known estimate over remaining budget before provider (Pylon #7691) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/contracts.json | 3 + tests/integration/proxy_config.yaml | 1 + .../integration/spend/test_cache_and_quota.py | 65 +++++++++++++++++++ 3 files changed, 69 insertions(+) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index d12e3ae4620..a3056d87077 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -107,6 +107,9 @@ "tests/integration/spend/test_cache_and_quota.py::test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores": [ "quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores" ], + "tests/integration/spend/test_cache_and_quota.py::test_fail_closed_key_rejects_known_estimate_over_remaining_budget_before_provider": [ + "quota_management.budget.key.fail_closed_rejects_estimate_over_remaining_headroom_before_provider" + ], "tests/integration/spend/test_cache_and_quota.py::test_different_system_messages_do_not_share_a_cached_response": [ "quota_management.response_cache.system_messages_partition_cache_identity" ], diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index a3b07f76d2f..08b36cc8397 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -5,6 +5,7 @@ general_settings: store_model_in_db: true disable_spend_logs: false proxy_batch_write_at: 1 + fail_closed_budget_enforcement: true litellm_settings: enable_redis_auth_cache: true cache: true diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index d32297765f6..5a0474331c0 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -1,3 +1,4 @@ +import re import uuid from contextlib import ExitStack from hashlib import sha256 @@ -211,6 +212,70 @@ def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gat assert upstream.get("/__observations").json()["requests"] == [] +@pytest.mark.covers("quota_management.budget.key.fail_closed_rejects_estimate_over_remaining_headroom_before_provider") +def test_fail_closed_key_rejects_known_estimate_over_remaining_budget_before_provider(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): + costly: Final = scenario.model(input_cost_per_token=0.005, output_cost_per_token=0.04) + cheap: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.001) + key: Final = scenario.key(models=[costly, cheap], max_budget=1.0) + digest: Final = sha256(key.encode()).hexdigest() + first: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": costly, "messages": [{"role": "user", "content": f"fill {uuid.uuid4().hex}"}], "max_tokens": 20}, + key=key, + ) + assert first.status_code == 200, first.text + assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(0.9) + spent: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.9, + seconds=70, + ) + assert float(spent[0]["spend"]) == pytest.approx(0.9) + upstream.get("/__observations").raise_for_status() + denied: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": costly, "messages": [{"role": "user", "content": f"over {uuid.uuid4().hex}"}], "max_tokens": 20}, + key=key, + ) + assert denied.status_code == 422, denied.text + error: Final = denied.json()["error"] + assert (error["type"], error["code"]) == ("budget_exceeded", "422"), denied.text + reported: Final = re.fullmatch( + r"Budget has been exceeded! Key=\w+ Current cost: (\S+), Estimated request cost: (\S+), Max budget: 1\.0", + error["message"], + ) + assert reported is not None, denied.text + assert float(reported.group(1)) == pytest.approx(0.9), denied.text + assert 20 * 0.04 < float(reported.group(2)) < 1.0, denied.text + assert upstream.get("/__observations").json()["requests"] == [] + assert float(gateway.get("/key/info", {"key": key})["info"]["spend"]) == pytest.approx(0.9) + fitting_prompt: Final = f"fits {uuid.uuid4().hex}" + served: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": cheap, "messages": [{"role": "user", "content": fitting_prompt}], "max_tokens": 5}, + key=key, + ) + assert served.status_code == 200, served.text + assert served.json()["usage"]["total_tokens"] == 40 + observed: Final = upstream.get("/__observations").json()["requests"] + assert [request["body"] for request in observed] == [ + {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": fitting_prompt}], "max_tokens": 5} + ], observed + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.94, + seconds=70, + ) + assert float(gateway.get("/key/info", {"key": key})["info"]["spend"]) == pytest.approx(0.94) + + @pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity") def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None: with (