test(integration): cover keys_auth endpoint contracts (#44674)

* test(integration): cover keys_auth endpoint contracts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): cover hashed, path and token spellings for keys_auth contracts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): strengthen JWT mapping lifecycle coverage

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): cover reset_spend over plaintext and hashed path spellings with reset_to 0 and 0.02

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): drop Final in reset loop and use drain-then-fresh-refusal in warmed-peer delete proof

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): disable model info refresh in JWT lifecycle rig

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): drop unused discovery constant

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): cover cache invalidation and key paths

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo <mateo@berri.ai>

tests/integration/authorization/test_jwt_key_mapping_lifecycle.py

@@ -0,0 +1,863 @@
+import json
+import os
+import time
+import uuid
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from hashlib import sha256
+from pathlib import Path
+from typing import Final, Literal
+
+import httpx
+import jwt
+import pytest
+from cryptography.hazmat.primitives.asymmetric import rsa
+from jwt.algorithms import RSAAlgorithm
+from pydantic import BaseModel, JsonValue, TypeAdapter
+from redis import Redis
+
+from tests.integration._support.client import (
+    Gateway,
+    delete_key_if_present,
+    eventually,
+    gateway_from_environment,
+)
+from tests.integration._support.database import read_rows
+from tests.integration._support.process import owned_proxy
+from tests.integration._support.wire import Reply, Request, Wire, wire_server
+
+_AUDIENCE: Final = "litellm-integration"
+_CLIENT_CLAIM: Final = "client_id"
+_KEY_ID: Final = "integration-signing-key"
+_UPSTREAM_KEY: Final = "synthetic-openai-key"
+_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
+_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
+pytestmark: Final = pytest.mark.timeout(180)
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class _AssistantMessage(BaseModel):
+    content: str
+
+
+class _Choice(BaseModel):
+    message: _AssistantMessage
+
+
+class _ChatResponse(BaseModel):
+    id: str
+    choices: tuple[_Choice, ...]
+    usage: _Usage
+
+
+class _KeyResponse(BaseModel):
+    key: str
+
+
+class _DeleteResponse(BaseModel):
+    deleted_keys: tuple[str, ...]
+
+
+class _JWTMappingResponse(BaseModel):
+    id: str
+    jwt_issuer: str | None
+    jwt_claim_name: str
+    jwt_claim_value: str
+    is_active: bool
+
+
+class _JWTMappingRow(BaseModel):
+    id: str
+    token: str
+    jwt_issuer: str
+    is_active: bool
+
+
+class _SpendRow(BaseModel):
+    api_key: str
+
+
+class _CountRow(BaseModel):
+    count: int
+
+
+class _MappingIdRow(BaseModel):
+    id: str
+
+
+class _ModelEntry(BaseModel):
+    id: str
+
+
+class _ModelList(BaseModel):
+    data: tuple[_ModelEntry, ...]
+
+
+class _UserMessage(BaseModel):
+    content: str
+
+
+class _ChatRequest(BaseModel):
+    model: str
+    messages: tuple[_UserMessage, ...]
+
+
+class _ProxyError(BaseModel):
+    message: str
+    type: str
+    param: str
+    code: str
+
+
+class _ErrorResponse(BaseModel):
+    error: _ProxyError
+
+
+@dataclass(frozen=True, slots=True)
+class _JWTLifecycleRuntime:
+    first: Gateway
+    second: Gateway
+    upstream: Wire
+    model: str
+    private_key_a: rsa.RSAPrivateKey
+    private_key_b: rsa.RSAPrivateKey
+    issuer_a: str
+    issuer_b: str
+
+
+def _proxy_config(
+    directory: Path,
+    model: str,
+    upstream_url: str,
+    issuers: tuple[dict[str, JsonValue], ...] = (),
+) -> Path:
+    config: Final = directory / "jwt_key_mapping_config.yaml"
+    config.write_text(
+        json.dumps(
+            {
+                "model_list": (
+                    {
+                        "model_name": model,
+                        "litellm_params": {
+                            "model": f"openai/{model}",
+                            "api_base": f"{upstream_url}/v1",
+                            "api_key": _UPSTREAM_KEY,
+                        },
+                    },
+                ),
+                "general_settings": {
+                    "master_key": "os.environ/LITELLM_MASTER_KEY",
+                    "database_url": "os.environ/DATABASE_URL",
+                    "store_model_in_db": True,
+                    "disable_model_info_refresh": True,
+                    "proxy_batch_write_at": 1,
+                    "proxy_batch_polling_interval": 1,
+                    "enable_jwt_auth": True,
+                    "litellm_jwtauth": {
+                        "user_id_jwt_field": "sub",
+                        "virtual_key_claim_field": _CLIENT_CLAIM,
+                        "unregistered_jwt_client_behavior": "reject",
+                        **({"issuers": issuers} if issuers else {}),
+                    },
+                },
+                "litellm_settings": {
+                    "enable_redis_auth_cache": True,
+                    "cache": True,
+                    "cache_params": {
+                        "type": "redis",
+                        "host": "os.environ/REDIS_HOST",
+                        "port": "os.environ/REDIS_PORT",
+                    },
+                },
+            }
+        )
+    )
+    return config
+
+
+def _issuer_config(issuer: str, jwks_url: str) -> dict[str, JsonValue]:
+    return {
+        "issuer": issuer,
+        "jwks_url": jwks_url,
+        "audience": _AUDIENCE,
+        "virtual_key_claim_field": _CLIENT_CLAIM,
+        "user_id_jwt_field": "sub",
+    }
+
+
+def _jwks_body(private_key: rsa.RSAPrivateKey) -> bytes:
+    public_jwk: Final = _CHAT_BODY.validate_json(RSAAlgorithm.to_jwk(private_key.public_key()))
+    return json.dumps({"keys": [{**public_jwk, "kid": _KEY_ID, "use": "sig", "alg": "RS256"}]}).encode()
+
+
+def _jwks_respond(private_key: rsa.RSAPrivateKey) -> Callable[[Request], Reply]:
+    body: Final = _jwks_body(private_key)
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "GET"
+        assert request.target == "/jwks"
+        return Reply(body=body)
+
+    return respond
+
+
+def _signed_jwt(private_key: rsa.RSAPrivateKey, client_id: str, issuer: str | None = None) -> str:
+    now: Final = int(time.time())
+    claims: Final[dict[str, JsonValue]] = {
+        "sub": "integration-user-" + uuid.uuid4().hex,
+        "client_id": client_id,
+        "aud": _AUDIENCE,
+        "iat": now,
+        "exp": now + 300,
+        **({"iss": issuer} if issuer is not None else {}),
+    }
+    return jwt.encode(claims, private_key, algorithm="RS256", headers={"kid": _KEY_ID})
+
+
+def _chat_body(model: str, prompt: str) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+
+
+def _wire_reply(identity: str) -> Reply:
+    return Reply(
+        body=json.dumps(
+            {
+                "id": identity,
+                "object": "chat.completion",
+                "created": 1,
+                "model": "integration-jwt-key-mapping",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "wire response"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
+            }
+        ).encode()
+    )
+
+
+def _chat_respond(backend: str) -> Callable[[Request], Reply]:
+    def respond(request: Request) -> Reply:
+        if request.method == "GET":
+            assert request.target == "/v1/models", request
+            assert request.headers["authorization"] == f"Bearer {_UPSTREAM_KEY}"
+            assert request.body == b""
+            return Reply(
+                body=json.dumps(
+                    {
+                        "object": "list",
+                        "data": [
+                            {
+                                "id": backend,
+                                "object": "model",
+                                "created": 0,
+                                "owned_by": "integration",
+                            }
+                        ],
+                    }
+                ).encode()
+            )
+        assert request.method == "POST", request
+        assert request.target == "/v1/chat/completions"
+        assert request.headers["authorization"] == f"Bearer {_UPSTREAM_KEY}"
+        body: Final = _CHAT_BODY.validate_python(json.loads(request.body))
+        parsed: Final = _ChatRequest.model_validate_json(request.body)
+        assert len(parsed.messages) == 1
+        assert body == _chat_body(backend, parsed.messages[0].content), body
+        return _wire_reply("chatcmpl-" + uuid.uuid4().hex)
+
+    return respond
+
+
+def _chat(gateway: Gateway, model: str, token: str, prompt: str) -> httpx.Response:
+    return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt), key=token)
+
+
+def _expected_chat_request(backend: str, prompt: str) -> tuple[str, str, dict[str, JsonValue]]:
+    body: Final = _chat_body(backend, prompt)
+    return "POST", "/v1/chat/completions", body
+
+
+def _upstream_request_key(
+    request: tuple[str, str, dict[str, JsonValue] | None],
+) -> tuple[str, str, str]:
+    return request[0], request[1], "" if request[2] is None else json.dumps(request[2], sort_keys=True)
+
+
+def _assert_upstream_requests(
+    upstream: Wire,
+    expected: tuple[tuple[str, str, dict[str, JsonValue] | None], ...],
+) -> None:
+    observed: Final = tuple(
+        sorted(
+            (
+                (
+                    request.method,
+                    request.target,
+                    None if not request.body else _CHAT_BODY.validate_python(json.loads(request.body)),
+                )
+                for request in upstream.drain()
+            ),
+            key=_upstream_request_key,
+        )
+    )
+    assert observed == tuple(sorted(expected, key=_upstream_request_key)), observed
+
+
+def _assert_chat(gateway: Gateway, model: str, token: str, prompt: str, key_hash: str) -> _ChatResponse:
+    response: Final = _chat(gateway, model, token, prompt)
+    assert response.status_code == 200, response.text
+    payload: Final = _ChatResponse.model_validate_json(response.content)
+    assert payload.choices[0].message.content == "wire response", response.text
+    assert payload.usage == _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40), response.text
+    spend: Final = TypeAdapter(tuple[_SpendRow, ...]).validate_python(
+        eventually(
+            lambda: read_rows(
+                'SELECT api_key FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+                (payload.id,),
+            ),
+            lambda rows: len(rows) == 1,
+            seconds=70,
+        )
+    )
+    assert spend == (_SpendRow(api_key=key_hash),)
+    return payload
+
+
+def _model_is_available(gateway: Gateway, model: str) -> bool:
+    response: Final = gateway.request("GET", "/v1/models")
+    if response.status_code != 200:
+        return False
+    payload: Final = _ModelList.model_validate_json(response.content)
+    return any(entry.id == model for entry in payload.data)
+
+
+def _create_key(gateway: Gateway, model: str) -> str:
+    response: Final = gateway.request("POST", "/key/generate", {"models": [model]})
+    assert response.status_code == 200, response.text
+    return _KeyResponse.model_validate_json(response.content).key
+
+
+def _create_mapping(
+    gateway: Gateway,
+    key: str,
+    client_id: str,
+    issuer: str | None = None,
+    spelling: Literal["key", "token"] = "key",
+) -> _JWTMappingResponse:
+    body: Final = {
+        "jwt_claim_name": _CLIENT_CLAIM,
+        "jwt_claim_value": client_id,
+        spelling: key if spelling == "key" else sha256(key.encode()).hexdigest(),
+        **({"jwt_issuer": issuer} if issuer is not None else {}),
+    }
+    response: Final = gateway.request("POST", "/jwt/key/mapping/new", body)
+    assert response.status_code == 200, response.text
+    return _JWTMappingResponse.model_validate_json(response.content)
+
+
+def _delete_mapping_if_present(gateway: Gateway, identity: str) -> None:
+    rows: Final = TypeAdapter(tuple[_MappingIdRow, ...]).validate_python(
+        read_rows('SELECT id FROM "LiteLLM_JWTKeyMapping" WHERE id=%s', (identity,))
+    )
+    if not rows:
+        return
+    response: Final = gateway.request("POST", "/jwt/key/mapping/delete", {"id": identity})
+    assert response.status_code == 200, response.text
+
+
+def _assert_unmapped(response: httpx.Response, client_id: str) -> None:
+    assert response.status_code == 403, response.text
+    error: Final = _ErrorResponse.model_validate_json(response.content)
+    expected: Final = f"JWT Key Mapping: No registered mapping for {_CLIENT_CLAIM}='{client_id}'. Access denied."
+    assert error.error == _ProxyError(message=expected, type="auth_error", param="None", code="403"), response.text
+
+
+def _mapping_info(gateway: Gateway, identity: str) -> _JWTMappingResponse:
+    response: Final = gateway.request("GET", "/jwt/key/mapping/info", params={"id": identity})
+    assert response.status_code == 200, response.text
+    return _JWTMappingResponse.model_validate_json(response.content)
+
+
+def _mapping_row(gateway: Gateway, identity: str) -> _JWTMappingRow:
+    rows: Final = TypeAdapter(tuple[_JWTMappingRow, ...]).validate_python(
+        read_rows(
+            'SELECT id, token, jwt_issuer, is_active FROM "LiteLLM_JWTKeyMapping" WHERE id=%s',
+            (identity,),
+        )
+    )
+    assert len(rows) == 1
+    return rows[0]
+
+
+def _spend_count(gateway: Gateway, key_hash: str) -> int:
+    rows: Final = TypeAdapter(tuple[_CountRow, ...]).validate_python(
+        read_rows('SELECT count(*) AS count FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (key_hash,))
+    )
+    assert len(rows) == 1
+    return rows[0].count
+
+
+@pytest.fixture(scope="module")
+def jwt_lifecycle(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_JWTLifecycleRuntime]:
+    directory: Final = tmp_path_factory.mktemp("jwt-key-mapping-lifecycle")
+    model: Final = "jwt-key-mapping-" + uuid.uuid4().hex
+    issuer_a: Final = "https://idp-a.integration.example"
+    issuer_b: Final = "https://idp-b.integration.example"
+    private_key_a: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    private_key_b: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+
+    with (
+        gateway_from_environment() as base,
+        wire_server(_chat_respond(model)) as upstream,
+        wire_server(_jwks_respond(private_key_a)) as jwks_a,
+        wire_server(_jwks_respond(private_key_b)) as jwks_b,
+        Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache,
+    ):
+        issuers: Final = (
+            _issuer_config(issuer_a, f"{jwks_a.url}/jwks"),
+            _issuer_config(issuer_b, f"{jwks_b.url}/jwks"),
+        )
+        config: Final = _proxy_config(directory, model, upstream.url, issuers)
+        baseline: Final = cache.pubsub_numsub(_CHANNEL)[0][1]
+        with (
+            owned_proxy(base, directory / "worker-one", {}, config=config) as first,
+            owned_proxy(base, directory / "worker-two", {}, config=config) as second,
+        ):
+            subscribers: Final = eventually(
+                lambda: cache.pubsub_numsub(_CHANNEL)[0][1],
+                lambda count: count >= baseline + 2,
+                seconds=20,
+            )
+            assert subscribers >= baseline + 2
+            eventually(lambda: _model_is_available(first, model), bool, seconds=30)
+            eventually(lambda: _model_is_available(second, model), bool, seconds=30)
+            _assert_upstream_requests(upstream, ())
+            yield _JWTLifecycleRuntime(
+                first=first,
+                second=second,
+                upstream=upstream,
+                model=model,
+                private_key_a=private_key_a,
+                private_key_b=private_key_b,
+                issuer_a=issuer_a,
+                issuer_b=issuer_b,
+            )
+
+
+def test_updating_a_warmed_mapping_deactivates_and_repoints_it_on_both_workers(
+    jwt_lifecycle: _JWTLifecycleRuntime,
+) -> None:
+    first: Final = jwt_lifecycle.first
+    second: Final = jwt_lifecycle.second
+    upstream: Final = jwt_lifecycle.upstream
+    model: Final = jwt_lifecycle.model
+    backend: Final = model
+    client_id: Final = "integration-client-" + uuid.uuid4().hex
+    issuer: Final = jwt_lifecycle.issuer_a
+    prompts: Final = (
+        "mapping warm primary " + uuid.uuid4().hex,
+        "mapping warm peer " + uuid.uuid4().hex,
+        "mapping inactive primary " + uuid.uuid4().hex,
+        "mapping inactive peer " + uuid.uuid4().hex,
+        "mapping repointed primary " + uuid.uuid4().hex,
+        "mapping repointed peer " + uuid.uuid4().hex,
+        "mapping restored primary " + uuid.uuid4().hex,
+        "mapping restored peer " + uuid.uuid4().hex,
+    )
+    token: Final = _signed_jwt(jwt_lifecycle.private_key_a, client_id, issuer)
+    upstream.drain()
+    with first.scenario() as scenario:
+        old_key: Final = _create_key(first, model)
+        new_key: Final = _create_key(first, model)
+        scenario.cleanups.callback(delete_key_if_present, first, old_key)
+        scenario.cleanups.callback(delete_key_if_present, first, new_key)
+        old_hash: Final = sha256(old_key.encode()).hexdigest()
+        new_hash: Final = sha256(new_key.encode()).hexdigest()
+        mapping: Final = _create_mapping(first, old_key, client_id, issuer)
+        scenario.cleanups.callback(_delete_mapping_if_present, first, mapping.id)
+        assert mapping.jwt_issuer == issuer
+        assert _mapping_row(first, mapping.id).jwt_issuer == issuer
+
+        first_response: Final = _assert_chat(first, model, token, prompts[0], old_hash)
+        second_response: Final = _assert_chat(second, model, token, prompts[1], old_hash)
+        assert first_response.id != second_response.id
+
+        deactivated_response: Final = first.request(
+            "POST", "/jwt/key/mapping/update", {"id": mapping.id, "is_active": False}
+        )
+        assert deactivated_response.status_code == 200, deactivated_response.text
+        deactivated: Final = _JWTMappingResponse.model_validate_json(deactivated_response.content)
+        assert deactivated.is_active is False, deactivated_response.text
+        assert _mapping_info(first, mapping.id).is_active is False
+        inactive_row: Final = _mapping_row(first, mapping.id)
+        assert inactive_row.token == old_hash and inactive_row.is_active is False
+        old_spend_count: Final = _spend_count(first, old_hash)
+
+        refused_first: Final = _chat(first, model, token, prompts[2])
+        _assert_unmapped(refused_first, client_id)
+        refused_second: Final = eventually(
+            lambda: _chat(second, model, token, prompts[3]),
+            lambda response: response.status_code == 403,
+            seconds=3,
+        )
+        _assert_unmapped(refused_second, client_id)
+        assert _spend_count(first, old_hash) == old_spend_count
+
+        reactivated_response: Final = first.request(
+            "POST",
+            "/jwt/key/mapping/update",
+            {"id": mapping.id, "is_active": True, "token": new_hash},
+        )
+        assert reactivated_response.status_code == 200, reactivated_response.text
+        reactivated: Final = _JWTMappingResponse.model_validate_json(reactivated_response.content)
+        assert reactivated.is_active is True, reactivated_response.text
+        assert _mapping_info(first, mapping.id).is_active is True
+        active_row: Final = _mapping_row(first, mapping.id)
+        assert active_row.token == new_hash and active_row.is_active is True
+        _assert_chat(first, model, token, prompts[4], new_hash)
+        _assert_chat(second, model, token, prompts[5], new_hash)
+
+        restored_response: Final = first.request(
+            "POST",
+            "/jwt/key/mapping/update",
+            {"id": mapping.id, "key": old_key},
+        )
+        assert restored_response.status_code == 200, restored_response.text
+        restored: Final = _JWTMappingResponse.model_validate_json(restored_response.content)
+        assert restored.is_active is True, restored_response.text
+        assert _mapping_info(first, mapping.id).is_active is True
+        restored_row: Final = _mapping_row(first, mapping.id)
+        assert restored_row.token == old_hash and restored_row.is_active is True
+        _assert_chat(first, model, token, prompts[6], old_hash)
+        _assert_chat(second, model, token, prompts[7], old_hash)
+        served_prompts: Final = prompts[0:2] + prompts[4:8]
+        _assert_upstream_requests(upstream, tuple(_expected_chat_request(backend, prompt) for prompt in served_prompts))
+
+
+def test_reactivating_a_global_mapping_restores_issuer_jwt_callers_on_both_workers(
+    jwt_lifecycle: _JWTLifecycleRuntime,
+) -> None:
+    pytest.skip(
+        "BUG: reactivating a global JWT key mapping keeps refusing issuer-scoped JWTs with 403 No registered mapping "
+        "because the miss was cached under the issuer key and update evicts only the global key"
+    )
+    first: Final = jwt_lifecycle.first
+    second: Final = jwt_lifecycle.second
+    upstream: Final = jwt_lifecycle.upstream
+    model: Final = jwt_lifecycle.model
+    client_id: Final = "integration-client-" + uuid.uuid4().hex
+    prompts: Final = tuple(f"global mapping reactivation {index} {uuid.uuid4().hex}" for index in range(6))
+    token: Final = _signed_jwt(jwt_lifecycle.private_key_a, client_id, jwt_lifecycle.issuer_a)
+    upstream.drain()
+
+    with first.scenario() as scenario:
+        key: Final = _create_key(first, model)
+        scenario.cleanups.callback(delete_key_if_present, first, key)
+        key_hash: Final = sha256(key.encode()).hexdigest()
+        mapping: Final = _create_mapping(first, key, client_id)
+        scenario.cleanups.callback(_delete_mapping_if_present, first, mapping.id)
+        assert mapping.jwt_issuer is None
+        assert _mapping_info(first, mapping.id).jwt_issuer is None
+        assert _mapping_row(first, mapping.id) == _JWTMappingRow(
+            id=mapping.id,
+            token=key_hash,
+            jwt_issuer="",
+            is_active=True,
+        )
+        _assert_chat(first, model, token, prompts[0], key_hash)
+        _assert_chat(second, model, token, prompts[1], key_hash)
+
+        deactivated_response: Final = first.request(
+            "POST",
+            "/jwt/key/mapping/update",
+            {"id": mapping.id, "is_active": False},
+        )
+        assert deactivated_response.status_code == 200, deactivated_response.text
+        assert _mapping_info(first, mapping.id).is_active is False
+        assert _mapping_row(first, mapping.id) == _JWTMappingRow(
+            id=mapping.id,
+            token=key_hash,
+            jwt_issuer="",
+            is_active=False,
+        )
+        _assert_unmapped(_chat(first, model, token, prompts[2]), client_id)
+        refused_second: Final = eventually(
+            lambda: _chat(second, model, token, prompts[3]),
+            lambda response: response.status_code == 403,
+            seconds=3,
+        )
+        _assert_unmapped(refused_second, client_id)
+        assert _spend_count(first, key_hash) == 2
+
+        reactivated_response: Final = first.request(
+            "POST",
+            "/jwt/key/mapping/update",
+            {"id": mapping.id, "is_active": True},
+        )
+        assert reactivated_response.status_code == 200, reactivated_response.text
+        reactivated: Final = _JWTMappingResponse.model_validate_json(reactivated_response.content)
+        assert reactivated.is_active is True, reactivated_response.text
+        assert reactivated.jwt_issuer is None, reactivated_response.text
+        assert _mapping_info(first, mapping.id).jwt_issuer is None
+        assert _mapping_row(first, mapping.id) == _JWTMappingRow(
+            id=mapping.id,
+            token=key_hash,
+            jwt_issuer="",
+            is_active=True,
+        )
+        _assert_chat(first, model, token, prompts[4], key_hash)
+        _assert_chat(second, model, token, prompts[5], key_hash)
+        assert _spend_count(first, key_hash) == 4
+        _assert_upstream_requests(
+            upstream,
+            tuple(_expected_chat_request(model, prompt) for prompt in prompts[0:2] + prompts[4:6]),
+        )
+
+
+@pytest.mark.parametrize("scope", ("global", "issuer"))
+@pytest.mark.parametrize("spelling", ("body", "path", "plaintext_path"))
+def test_regenerating_a_jwt_mapped_key_moves_jwt_callers_to_the_new_hash(
+    jwt_lifecycle: _JWTLifecycleRuntime, spelling: str, scope: str
+) -> None:
+    first: Final = jwt_lifecycle.first
+    second: Final = jwt_lifecycle.second
+    upstream: Final = jwt_lifecycle.upstream
+    model: Final = jwt_lifecycle.model
+    backend: Final = model
+    client_id: Final = "integration-client-" + uuid.uuid4().hex
+    prompts: Final = (
+        "regenerate warm primary " + uuid.uuid4().hex,
+        "regenerate warm peer " + uuid.uuid4().hex,
+        "regenerated primary " + uuid.uuid4().hex,
+        "regenerated peer " + uuid.uuid4().hex,
+    )
+    issuer: Final = jwt_lifecycle.issuer_a if scope == "issuer" else None
+    token: Final = _signed_jwt(jwt_lifecycle.private_key_a, client_id, jwt_lifecycle.issuer_a)
+    upstream.drain()
+    with first.scenario() as scenario:
+        old_key: Final = _create_key(first, model)
+        scenario.cleanups.callback(delete_key_if_present, first, old_key)
+        old_hash: Final = sha256(old_key.encode()).hexdigest()
+        mapping: Final = _create_mapping(first, old_key, client_id, issuer)
+        scenario.cleanups.callback(_delete_mapping_if_present, first, mapping.id)
+        expected_issuer: Final = issuer if issuer is not None else ""
+        assert mapping.jwt_issuer == issuer
+        assert _mapping_info(first, mapping.id).jwt_issuer == issuer
+        assert _mapping_row(first, mapping.id).jwt_issuer == expected_issuer
+        _assert_chat(first, model, token, prompts[0], old_hash)
+        _assert_chat(second, model, token, prompts[1], old_hash)
+
+        path_key: Final = old_hash if spelling == "path" else old_key
+        regenerated_response: Final = (
+            first.request("POST", "/key/regenerate", {"key": old_key})
+            if spelling == "body"
+            else first.request("POST", f"/key/{path_key}/regenerate", {})
+        )
+        assert regenerated_response.status_code == 200, regenerated_response.text
+        new_key: Final = _KeyResponse.model_validate_json(regenerated_response.content).key
+        scenario.cleanups.callback(delete_key_if_present, first, new_key)
+        new_hash: Final = sha256(new_key.encode()).hexdigest()
+        mapped_row: Final = _mapping_row(first, mapping.id)
+        assert mapped_row.token == new_hash and mapped_row.is_active is True
+        assert mapped_row.jwt_issuer == expected_issuer
+        assert _mapping_info(first, mapping.id).jwt_issuer == issuer
+        _assert_chat(first, model, token, prompts[2], new_hash)
+        _assert_chat(second, model, token, prompts[3], new_hash)
+        _assert_upstream_requests(upstream, tuple(_expected_chat_request(backend, prompt) for prompt in prompts))
+
+
+@pytest.mark.parametrize("scope", ("global", "issuer"))
+def test_deleting_a_jwt_mapped_key_stops_its_mapping_on_both_workers(
+    jwt_lifecycle: _JWTLifecycleRuntime, scope: str
+) -> None:
+    first: Final = jwt_lifecycle.first
+    second: Final = jwt_lifecycle.second
+    upstream: Final = jwt_lifecycle.upstream
+    model: Final = jwt_lifecycle.model
+    backend: Final = model
+    client_id: Final = "integration-client-" + uuid.uuid4().hex
+    prompts: Final = (
+        "delete mapped warm primary " + uuid.uuid4().hex,
+        "delete mapped warm peer " + uuid.uuid4().hex,
+        "delete mapped primary " + uuid.uuid4().hex,
+        "delete mapped peer " + uuid.uuid4().hex,
+    )
+    issuer: Final = jwt_lifecycle.issuer_a if scope == "issuer" else None
+    token: Final = _signed_jwt(jwt_lifecycle.private_key_a, client_id, jwt_lifecycle.issuer_a)
+    upstream.drain()
+    with first.scenario() as scenario:
+        key: Final = _create_key(first, model)
+        scenario.cleanups.callback(delete_key_if_present, first, key)
+        key_hash: Final = sha256(key.encode()).hexdigest()
+        mapping: Final = _create_mapping(first, key, client_id, issuer)
+        scenario.cleanups.callback(_delete_mapping_if_present, first, mapping.id)
+        expected_issuer: Final = issuer if issuer is not None else ""
+        assert mapping.jwt_issuer == issuer
+        assert _mapping_info(first, mapping.id).jwt_issuer == issuer
+        assert _mapping_row(first, mapping.id).jwt_issuer == expected_issuer
+        warm_first: Final = _assert_chat(first, model, token, prompts[0], key_hash)
+        warm_second: Final = _assert_chat(second, model, token, prompts[1], key_hash)
+        assert warm_first.id != warm_second.id
+        warm_spend_count: Final = _spend_count(first, key_hash)
+        assert warm_spend_count == 2
+
+        deletion: Final = first.request("POST", "/key/delete", {"keys": [key]})
+        assert deletion.status_code == 200, deletion.text
+        assert _DeleteResponse.model_validate_json(deletion.content).deleted_keys == (key,), deletion.text
+        assert read_rows('SELECT id FROM "LiteLLM_JWTKeyMapping" WHERE id=%s', (mapping.id,)) == []
+        second_refusal: Final = eventually(
+            lambda: _chat(second, model, token, prompts[2]),
+            lambda response: response.status_code == 403,
+            seconds=3,
+        )
+        _assert_unmapped(second_refusal, client_id)
+        first_refusal: Final = _chat(first, model, token, prompts[3])
+        _assert_unmapped(first_refusal, client_id)
+        assert _spend_count(first, key_hash) == warm_spend_count
+        _assert_upstream_requests(upstream, tuple(_expected_chat_request(backend, prompt) for prompt in prompts[:2]))
+
+
+def test_issuer_scoped_mappings_with_the_same_claim_resolve_per_issuer(
+    jwt_lifecycle: _JWTLifecycleRuntime,
+) -> None:
+    candidate: Final = jwt_lifecycle.first
+    upstream: Final = jwt_lifecycle.upstream
+    model: Final = jwt_lifecycle.model
+    backend: Final = model
+    client_id: Final = "integration-client-" + uuid.uuid4().hex
+    issuer_a: Final = jwt_lifecycle.issuer_a
+    issuer_b: Final = jwt_lifecycle.issuer_b
+    prompts: Final = tuple(f"issuer mapping {index} {uuid.uuid4().hex}" for index in range(4))
+    upstream.drain()
+
+    with candidate.scenario() as scenario:
+        key_a: Final = _create_key(candidate, model)
+        key_b: Final = _create_key(candidate, model)
+        scenario.cleanups.callback(delete_key_if_present, candidate, key_a)
+        scenario.cleanups.callback(delete_key_if_present, candidate, key_b)
+        hash_a: Final = sha256(key_a.encode()).hexdigest()
+        hash_b: Final = sha256(key_b.encode()).hexdigest()
+        mapping_a: Final = _create_mapping(candidate, key_a, client_id, issuer_a)
+        mapping_b: Final = _create_mapping(candidate, key_b, client_id, issuer_b, spelling="token")
+        scenario.cleanups.callback(_delete_mapping_if_present, candidate, mapping_a.id)
+        scenario.cleanups.callback(_delete_mapping_if_present, candidate, mapping_b.id)
+
+        _assert_chat(
+            candidate,
+            model,
+            _signed_jwt(jwt_lifecycle.private_key_a, client_id, issuer_a),
+            prompts[0],
+            hash_a,
+        )
+        _assert_chat(
+            candidate,
+            model,
+            _signed_jwt(jwt_lifecycle.private_key_b, client_id, issuer_b),
+            prompts[1],
+            hash_b,
+        )
+        _assert_chat(
+            candidate,
+            model,
+            _signed_jwt(jwt_lifecycle.private_key_a, client_id, issuer_a),
+            prompts[2],
+            hash_a,
+        )
+        _assert_chat(
+            candidate,
+            model,
+            _signed_jwt(jwt_lifecycle.private_key_b, client_id, issuer_b),
+            prompts[3],
+            hash_b,
+        )
+        rows: Final = TypeAdapter(tuple[_JWTMappingRow, ...]).validate_python(
+            read_rows(
+                'SELECT id, token, jwt_issuer, is_active FROM "LiteLLM_JWTKeyMapping" '
+                "WHERE jwt_claim_name=%s AND jwt_claim_value=%s ORDER BY jwt_issuer",
+                (_CLIENT_CLAIM, client_id),
+            )
+        )
+        assert rows == (
+            _JWTMappingRow(id=mapping_a.id, token=hash_a, jwt_issuer=issuer_a, is_active=True),
+            _JWTMappingRow(id=mapping_b.id, token=hash_b, jwt_issuer=issuer_b, is_active=True),
+        )
+        _assert_upstream_requests(
+            upstream,
+            tuple(_expected_chat_request(backend, prompt) for prompt in prompts),
+        )
+
+
+def test_repointing_a_warmed_mapping_to_a_new_issuer_moves_it_on_both_workers(
+    jwt_lifecycle: _JWTLifecycleRuntime,
+) -> None:
+    first: Final = jwt_lifecycle.first
+    second: Final = jwt_lifecycle.second
+    upstream: Final = jwt_lifecycle.upstream
+    model: Final = jwt_lifecycle.model
+    issuer_a: Final = jwt_lifecycle.issuer_a
+    issuer_b: Final = jwt_lifecycle.issuer_b
+    client_id: Final = "integration-client-" + uuid.uuid4().hex
+    prompts: Final = tuple(f"mapping repoint {index} {uuid.uuid4().hex}" for index in range(4))
+    upstream.drain()
+
+    with first.scenario() as scenario:
+        key: Final = _create_key(first, model)
+        scenario.cleanups.callback(delete_key_if_present, first, key)
+        key_hash: Final = sha256(key.encode()).hexdigest()
+        mapping: Final = _create_mapping(first, key, client_id, issuer_a)
+        scenario.cleanups.callback(_delete_mapping_if_present, first, mapping.id)
+        assert _mapping_row(first, mapping.id) == _JWTMappingRow(
+            id=mapping.id,
+            token=key_hash,
+            jwt_issuer=issuer_a,
+            is_active=True,
+        )
+
+        new_identity_token: Final = _signed_jwt(jwt_lifecycle.private_key_b, client_id, issuer_b)
+        _assert_unmapped(_chat(first, model, new_identity_token, prompts[0]), client_id)
+        _assert_unmapped(_chat(second, model, new_identity_token, prompts[1]), client_id)
+
+        old_token: Final = _signed_jwt(jwt_lifecycle.private_key_a, client_id, issuer_a)
+        _assert_chat(first, model, old_token, prompts[2], key_hash)
+        _assert_chat(second, model, old_token, prompts[3], key_hash)
+        assert _spend_count(first, key_hash) == 2
+
+        update: Final = first.request(
+            "POST",
+            "/jwt/key/mapping/update",
+            {"id": mapping.id, "jwt_issuer": issuer_b},
+        )
+        assert update.status_code == 200, update.text
+        updated: Final = _JWTMappingResponse.model_validate_json(update.content)
+        assert updated.jwt_claim_value == client_id, update.text
+        assert updated.jwt_issuer == issuer_b, update.text
+        assert _mapping_info(first, mapping.id) == updated
+        assert _mapping_row(first, mapping.id) == _JWTMappingRow(
+            id=mapping.id,
+            token=key_hash,
+            jwt_issuer=issuer_b,
+            is_active=True,
+        )
+
+        _assert_chat(first, model, new_identity_token, prompts[0], key_hash)
+        _assert_chat(second, model, new_identity_token, prompts[1], key_hash)
+        assert _spend_count(first, key_hash) == 4
+        _assert_unmapped(_chat(first, model, old_token, prompts[2]), client_id)
+        _assert_unmapped(_chat(second, model, old_token, prompts[3]), client_id)
+        assert _spend_count(first, key_hash) == 4
+        _assert_upstream_requests(
+            upstream,
+            tuple(_expected_chat_request(model, prompt) for prompt in prompts),
+        )

tests/integration/authorization/test_warmed_policy.py

@@ -1,3 +1,4 @@
+import json
 import os
 import uuid
 from contextlib import ExitStack
@@ -9,11 +10,187 @@ import psycopg
 import pytest
 from hypothesis import strategies as st
 from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
-from pydantic import BaseModel, ConfigDict, JsonValue
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
 
 from tests.integration._support.client import Gateway, eventually, object_value, string_value, team_admin_permissions
 from tests.integration._support.database import read_rows
 from tests.integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
+from tests.integration._support.wire import Reply, Request, wire_server
+
+_CHAT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
+_EXPECTED_CHAT: Final = ("POST", "/v1/chat/completions")
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class _AssistantMessage(BaseModel):
+    content: str
+
+
+class _Choice(BaseModel):
+    message: _AssistantMessage
+
+
+class _ChatResponse(BaseModel):
+    id: str
+    choices: tuple[_Choice, ...]
+    usage: _Usage
+
+
+class _ModelEntry(BaseModel):
+    id: str
+
+
+class _ModelList(BaseModel):
+    data: tuple[_ModelEntry, ...]
+
+
+class _KeyInfo(BaseModel):
+    blocked: bool
+
+
+class _KeyInfoResponse(BaseModel):
+    info: _KeyInfo
+
+
+class _BlockedRow(BaseModel):
+    blocked: bool
+
+
+class _AuthError(BaseModel):
+    type: str
+
+
+class _ErrorResponse(BaseModel):
+    error: _AuthError
+
+
+def _chat_body(model: str, prompt: str) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+
+
+def _wire_completion(identity: str) -> Reply:
+    return Reply(
+        body=json.dumps(
+            {
+                "id": identity,
+                "object": "chat.completion",
+                "created": 1,
+                "model": "integration-keys-auth-block",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "wire response"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
+            }
+        ).encode()
+    )
+
+
+def _wire_chat(gateway: Gateway, model: str, key: str, prompt: str) -> httpx.Response:
+    return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt), key=key)
+
+
+def _assert_wire_success(response: httpx.Response) -> _ChatResponse:
+    assert response.status_code == 200, response.text
+    payload: Final = _ChatResponse.model_validate_json(response.content)
+    assert payload.choices[0].message.content == "wire response", response.text
+    assert payload.usage == _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40), response.text
+    return payload
+
+
+def _model_is_available(gateway: Gateway, model: str) -> bool:
+    response: Final = gateway.request("GET", "/v1/models")
+    if response.status_code != 200:
+        return False
+    models: Final = _ModelList.model_validate_json(response.content)
+    return model in tuple(entry.id for entry in models.data)
+
+
+@pytest.mark.parametrize("hashed", (False, True), ids=("plaintext", "hashed"))
+def test_key_block_and_unblock_reach_both_warmed_workers(gateway: Gateway, peer: Gateway, hashed: bool) -> None:
+    backend: Final = "integration-keys-auth-block"
+    prompts: Final = (
+        "block warm primary " + uuid.uuid4().hex,
+        "block warm peer " + uuid.uuid4().hex,
+        "block rejected primary " + uuid.uuid4().hex,
+        "block rejected peer " + uuid.uuid4().hex,
+        "unblock primary " + uuid.uuid4().hex,
+        "unblock peer " + uuid.uuid4().hex,
+    )
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.target == "/v1/chat/completions"
+        assert request.headers["authorization"] == "Bearer synthetic-openai-key"
+        body: Final = _CHAT_REQUEST.validate_json(request.body)
+        assert body in tuple(_chat_body(backend, prompt) for prompt in prompts), body
+        return _wire_completion("chatcmpl-" + uuid.uuid4().hex)
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model=f"openai/{backend}", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
+        )
+        eventually(lambda: _model_is_available(peer, model), bool, seconds=30)
+        key: Final = scenario.key(models=[model])
+        key_hash: Final = sha256(key.encode()).hexdigest()
+        endpoint_key: Final = key_hash if hashed else key
+        warmed: Final = tuple(
+            _assert_wire_success(_wire_chat(worker, model, key, prompt))
+            for worker, prompt in zip((gateway, peer), prompts[:2], strict=True)
+        )
+        assert len(warmed) == 2
+        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2
+
+        blocked: Final = gateway.request("POST", "/key/block", {"key": endpoint_key})
+        assert blocked.status_code == 200, blocked.text
+        assert _KeyInfo.model_validate_json(blocked.content).blocked is True, blocked.text
+        primary_refusal: Final = _wire_chat(gateway, model, key, prompts[2])
+        assert primary_refusal.status_code == 401, primary_refusal.text
+        assert _ErrorResponse.model_validate_json(primary_refusal.content).error.type == "auth_error", (
+            primary_refusal.text
+        )
+        peer_refusal: Final = eventually(
+            lambda: _wire_chat(peer, model, key, prompts[3]),
+            lambda response: response.status_code == 401,
+            seconds=3,
+        )
+        assert _ErrorResponse.model_validate_json(peer_refusal.content).error.type == "auth_error", peer_refusal.text
+        assert wire.drain() == (), "blocked requests reached the upstream"
+        blocked_info: Final = gateway.request("GET", "/key/info", params={"key": key})
+        assert blocked_info.status_code == 200, blocked_info.text
+        assert _KeyInfoResponse.model_validate_json(blocked_info.content).info.blocked is True, blocked_info.text
+        blocked_rows: Final = TypeAdapter(tuple[_BlockedRow, ...]).validate_python(
+            read_rows('SELECT blocked FROM "LiteLLM_VerificationToken" WHERE token=%s', (key_hash,))
+        )
+        assert blocked_rows == (_BlockedRow(blocked=True),)
+
+        unblocked: Final = gateway.request("POST", "/key/unblock", {"key": endpoint_key})
+        assert unblocked.status_code == 200, unblocked.text
+        assert _KeyInfo.model_validate_json(unblocked.content).blocked is False, unblocked.text
+        served: Final = tuple(
+            _assert_wire_success(_wire_chat(worker, model, key, prompt))
+            for worker, prompt in zip((gateway, peer), prompts[4:], strict=True)
+        )
+        assert len(served) == 2
+        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2
+        unblocked_info: Final = gateway.request("GET", "/key/info", params={"key": key})
+        assert unblocked_info.status_code == 200, unblocked_info.text
+        assert _KeyInfoResponse.model_validate_json(unblocked_info.content).info.blocked is False, unblocked_info.text
+        unblocked_rows: Final = TypeAdapter(tuple[_BlockedRow, ...]).validate_python(
+            read_rows('SELECT blocked FROM "LiteLLM_VerificationToken" WHERE token=%s', (key_hash,))
+        )
+        assert unblocked_rows == (_BlockedRow(blocked=False),)
 
 
 SCIM_CORE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:core:2.0:User"

tests/integration/database/test_reader_writer_regeneration.py

@@ -1,18 +1,21 @@
+import json
 import os
 import uuid
+from concurrent.futures import ThreadPoolExecutor
 from hashlib import sha256
 from pathlib import Path
-from concurrent.futures import ThreadPoolExecutor
 from typing import Final
 from urllib.parse import urlsplit, urlunsplit
 
+import httpx
 import psycopg
 import pytest
-from psycopg import sql
-
 from integration._support.client import Gateway, delete_key_if_present, eventually, string_value
 from integration._support.database import read_rows
 from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, wire_server
+from psycopg import sql
+from pydantic import BaseModel, JsonValue, TypeAdapter
 
 
 @pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants")
@@ -129,3 +132,207 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew
             admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
             admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
     assert read_rows("SELECT rolname FROM pg_roles WHERE rolname=%s", (role,)) == []
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class _AssistantMessage(BaseModel):
+    content: str
+
+
+class _Choice(BaseModel):
+    message: _AssistantMessage
+
+
+class _ChatResponse(BaseModel):
+    id: str
+    choices: tuple[_Choice, ...]
+    usage: _Usage
+
+
+class _GenerateKeyResponse(BaseModel):
+    key: str
+
+
+class _Error(BaseModel):
+    type: str
+
+
+class _ErrorResponse(BaseModel):
+    error: _Error
+
+
+class _SpendLog(BaseModel):
+    api_key: str
+    spend: float
+
+
+class _KeySpend(BaseModel):
+    spend: float
+
+
+class _DeprecatedKey(BaseModel):
+    active_token_id: str
+    grace_seconds: float
+
+
+_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
+_EXPECTED_CHAT: Final = ("POST", "/v1/chat/completions")
+
+
+def _chat_body(model: str, prompt: str) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+
+
+def _wire_reply(identity: str) -> Reply:
+    return Reply(
+        body=json.dumps(
+            {
+                "id": identity,
+                "object": "chat.completion",
+                "created": 1,
+                "model": "integration-keys-auth-grace",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "wire response"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
+            }
+        ).encode()
+    )
+
+
+def _chat(gateway: Gateway, model: str, key: str, prompt: str) -> httpx.Response:
+    return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt), key=key)
+
+
+def _assert_success(response: httpx.Response) -> _ChatResponse:
+    assert response.status_code == 200, response.text
+    payload: Final = _ChatResponse.model_validate_json(response.content)
+    assert payload.choices[0].message.content == "wire response", response.text
+    assert payload.usage == _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40), response.text
+    return payload
+
+
+@pytest.mark.parametrize("spelling", ("body", "path", "plaintext_path"))
+def test_regenerate_with_a_grace_period_keeps_the_old_key_serving_billed_to_the_new_row(
+    gateway: Gateway, spelling: str
+) -> None:
+    assert os.environ.get("LITELLM_KEY_ROTATION_GRACE_PERIOD") is None
+    backend: Final = "integration-keys-auth-grace"
+    prompts: Final = (
+        "grace old key " + uuid.uuid4().hex,
+        "grace new key " + uuid.uuid4().hex,
+        "revoked second key " + uuid.uuid4().hex,
+        "final generated key " + uuid.uuid4().hex,
+    )
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.target == "/v1/chat/completions"
+        assert request.headers["authorization"] == "Bearer synthetic-openai-key"
+        body: Final = _CHAT_BODY.validate_json(request.body)
+        assert body in tuple(_chat_body(backend, prompt) for prompt in prompts), body
+        return _wire_reply("chatcmpl-" + uuid.uuid4().hex)
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model=f"openai/{backend}",
+            api_base=f"{wire.url}/v1",
+            api_key="synthetic-openai-key",
+            input_cost_per_token=0.001,
+            output_cost_per_token=0.002,
+        )
+        created: Final = gateway.request("POST", "/key/generate", {"models": [model]})
+        assert created.status_code == 200, created.text
+        key_one: Final = _GenerateKeyResponse.model_validate_json(created.content).key
+        scenario.cleanups.callback(delete_key_if_present, gateway, key_one)
+        hash_one: Final = sha256(key_one.encode()).hexdigest()
+
+        path_key_one: Final = hash_one if spelling == "path" else key_one
+        generated_two: Final = (
+            gateway.request("POST", "/key/regenerate", {"key": key_one, "grace_period": "1h"})
+            if spelling == "body"
+            else gateway.request("POST", f"/key/{path_key_one}/regenerate", {"grace_period": "1h"})
+        )
+        assert generated_two.status_code == 200, generated_two.text
+        key_two: Final = _GenerateKeyResponse.model_validate_json(generated_two.content).key
+        scenario.cleanups.callback(delete_key_if_present, gateway, key_two)
+        hash_two: Final = sha256(key_two.encode()).hexdigest()
+        old_response: Final = _assert_success(_chat(gateway, model, key_one, prompts[0]))
+        new_response: Final = _assert_success(_chat(gateway, model, key_two, prompts[1]))
+        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2
+
+        old_spend: Final = TypeAdapter(tuple[_SpendLog, ...]).validate_python(
+            eventually(
+                lambda: read_rows(
+                    'SELECT api_key, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+                    (old_response.id,),
+                ),
+                lambda rows: len(rows) == 1,
+                seconds=70,
+            )
+        )
+        new_spend: Final = TypeAdapter(tuple[_SpendLog, ...]).validate_python(
+            eventually(
+                lambda: read_rows(
+                    'SELECT api_key, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+                    (new_response.id,),
+                ),
+                lambda rows: len(rows) == 1,
+                seconds=70,
+            )
+        )
+        assert old_spend[0].api_key == hash_two
+        assert old_spend[0].spend > 0
+        assert new_spend[0].api_key == hash_two
+        assert new_spend[0].spend > 0
+        key_spend: Final = TypeAdapter(tuple[_KeySpend, ...]).validate_python(
+            eventually(
+                lambda: read_rows(
+                    'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s',
+                    (hash_two,),
+                ),
+                lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
+                seconds=70,
+            )
+        )
+        assert key_spend[0].spend > 0
+        deprecated: Final = TypeAdapter(tuple[_DeprecatedKey, ...]).validate_python(
+            read_rows(
+                "SELECT active_token_id, EXTRACT(EPOCH FROM revoke_at - now())::float8 AS grace_seconds "
+                'FROM "LiteLLM_DeprecatedVerificationToken" WHERE token=%s',
+                (hash_one,),
+            )
+        )
+        assert len(deprecated) == 1
+        assert deprecated[0].active_token_id == hash_two
+        assert deprecated[0].grace_seconds == pytest.approx(3600, abs=120)
+
+        path_key_two: Final = hash_two if spelling == "path" else key_two
+        generated_three: Final = (
+            gateway.request("POST", "/key/regenerate", {"key": key_two})
+            if spelling == "body"
+            else gateway.request("POST", f"/key/{path_key_two}/regenerate", {})
+        )
+        assert generated_three.status_code == 200, generated_three.text
+        key_three: Final = _GenerateKeyResponse.model_validate_json(generated_three.content).key
+        scenario.cleanups.callback(delete_key_if_present, gateway, key_three)
+        refused: Final = _chat(gateway, model, key_two, prompts[2])
+        assert refused.status_code == 401, refused.text
+        assert _ErrorResponse.model_validate_json(refused.content).error.type == "token_not_found_in_db", refused.text
+        assert wire.drain() == (), "the revoked second key reached the upstream"
+        assert (
+            _assert_success(_chat(gateway, model, key_three, prompts[3])).choices[0].message.content == "wire response"
+        )
+        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,)

tests/integration/management/test_key_delete_revocation.py

@@ -0,0 +1,233 @@
+import json
+import uuid
+from hashlib import sha256
+from typing import Final
+
+import httpx
+import pytest
+from pydantic import BaseModel, JsonValue, TypeAdapter
+
+from tests.integration._support.client import Gateway, delete_key_if_present, eventually
+from tests.integration._support.database import read_rows
+from tests.integration._support.wire import Reply, Request, wire_server
+
+_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
+_EXPECTED_CHAT: Final = ("POST", "/v1/chat/completions")
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class _AssistantMessage(BaseModel):
+    content: str
+
+
+class _Choice(BaseModel):
+    message: _AssistantMessage
+
+
+class _ChatResponse(BaseModel):
+    id: str
+    choices: tuple[_Choice, ...]
+    usage: _Usage
+
+
+class _ModelEntry(BaseModel):
+    id: str
+
+
+class _ModelList(BaseModel):
+    data: tuple[_ModelEntry, ...]
+
+
+class _GenerateKeyResponse(BaseModel):
+    key: str
+
+
+class _DeleteResponse(BaseModel):
+    deleted_keys: tuple[str, ...]
+
+
+class _KeyInfo(BaseModel):
+    status: str
+
+
+class _KeyInfoResponse(BaseModel):
+    info: _KeyInfo
+
+
+class _Error(BaseModel):
+    type: str
+
+
+class _ErrorResponse(BaseModel):
+    error: _Error
+
+
+class _DeletedKeyRow(BaseModel):
+    token: str
+
+
+def _chat_body(model: str, prompt: str) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+
+
+def _wire_reply(identity: str) -> Reply:
+    return Reply(
+        body=json.dumps(
+            {
+                "id": identity,
+                "object": "chat.completion",
+                "created": 1,
+                "model": "integration-keys-auth-delete",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "wire response"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
+            }
+        ).encode()
+    )
+
+
+def _chat(gateway: Gateway, model: str, key: str, prompt: str) -> httpx.Response:
+    return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt), key=key)
+
+
+def _assert_success(response: httpx.Response) -> _ChatResponse:
+    assert response.status_code == 200, response.text
+    payload: Final = _ChatResponse.model_validate_json(response.content)
+    assert payload.choices[0].message.content == "wire response", response.text
+    assert payload.usage == _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40), response.text
+    return payload
+
+
+def _model_is_available(gateway: Gateway, model: str) -> bool:
+    response: Final = gateway.request("GET", "/v1/models")
+    if response.status_code != 200:
+        return False
+    models: Final = _ModelList.model_validate_json(response.content)
+    return model in tuple(entry.id for entry in models.data)
+
+
+@pytest.mark.parametrize("spelling", ("keys", "hashed_keys", "key_aliases"))
+def test_deleting_a_key_revokes_it(gateway: Gateway, spelling: str) -> None:
+    backend: Final = "integration-keys-auth-delete"
+    prompts: Final = (
+        "delete warm " + uuid.uuid4().hex,
+        "delete reject " + uuid.uuid4().hex,
+    )
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.target == "/v1/chat/completions"
+        assert request.headers["authorization"] == "Bearer synthetic-openai-key"
+        body: Final = _CHAT_BODY.validate_json(request.body)
+        assert body in tuple(_chat_body(backend, prompt) for prompt in prompts), body
+        return _wire_reply("chatcmpl-" + uuid.uuid4().hex)
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model=f"openai/{backend}", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
+        )
+        eventually(lambda: _model_is_available(gateway, model), bool, seconds=30)
+        alias: Final = "integration-delete-" + uuid.uuid4().hex
+        created_response: Final = gateway.request("POST", "/key/generate", {"key_alias": alias, "models": [model]})
+        assert created_response.status_code == 200, created_response.text
+        key: Final = _GenerateKeyResponse.model_validate_json(created_response.content).key
+        scenario.cleanups.callback(delete_key_if_present, gateway, key)
+        key_hash: Final = sha256(key.encode()).hexdigest()
+        _assert_success(_chat(gateway, model, key, prompts[0]))
+        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,)
+
+        identifier: Final = {"keys": key, "hashed_keys": key_hash, "key_aliases": alias}[spelling]
+        delete_body: Final = {"key_aliases" if spelling == "key_aliases" else "keys": [identifier]}
+        deleted: Final = gateway.request("POST", "/key/delete", delete_body)
+        assert deleted.status_code == 200, deleted.text
+        assert _DeleteResponse.model_validate_json(deleted.content).deleted_keys == (identifier,), deleted.text
+        primary_refusal: Final = _chat(gateway, model, key, prompts[1])
+        assert primary_refusal.status_code == 401, primary_refusal.text
+        assert _ErrorResponse.model_validate_json(primary_refusal.content).error.type == "token_not_found_in_db", (
+            primary_refusal.text
+        )
+        assert wire.drain() == (), "deleted-key requests reached the upstream"
+        info: Final = gateway.request("GET", "/key/info", params={"key": key_hash})
+        assert info.status_code == 200, info.text
+        assert _KeyInfoResponse.model_validate_json(info.content).info.status == "deleted", info.text
+        live_rows: Final = read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (key_hash,))
+        deleted_rows: Final = TypeAdapter(tuple[_DeletedKeyRow, ...]).validate_python(
+            read_rows('SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE token=%s', (key_hash,))
+        )
+        assert live_rows == []
+        assert deleted_rows == (_DeletedKeyRow(token=key_hash),)
+
+
+def test_deleting_a_key_revokes_it_on_a_warmed_peer(gateway: Gateway, peer: Gateway) -> None:
+    pytest.skip(
+        "BUG: POST /key/delete does not broadcast the key eviction, so a warmed peer keeps answering 200 for about "
+        "60 seconds until its in-memory cache TTL expires"
+    )
+    backend: Final = "integration-keys-auth-delete"
+    prompts: Final = (
+        "delete warm primary " + uuid.uuid4().hex,
+        "delete warm peer " + uuid.uuid4().hex,
+        "delete reject peer " + uuid.uuid4().hex,
+        "delete reject peer fresh " + uuid.uuid4().hex,
+    )
+
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.target == "/v1/chat/completions"
+        assert request.headers["authorization"] == "Bearer synthetic-openai-key"
+        body: Final = _CHAT_BODY.validate_json(request.body)
+        assert body in tuple(_chat_body(backend, prompt) for prompt in prompts), body
+        return _wire_reply("chatcmpl-" + uuid.uuid4().hex)
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model=f"openai/{backend}", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key"
+        )
+        eventually(lambda: _model_is_available(peer, model), bool, seconds=30)
+        alias: Final = "integration-delete-" + uuid.uuid4().hex
+        created_response: Final = gateway.request("POST", "/key/generate", {"key_alias": alias, "models": [model]})
+        assert created_response.status_code == 200, created_response.text
+        key: Final = _GenerateKeyResponse.model_validate_json(created_response.content).key
+        scenario.cleanups.callback(delete_key_if_present, gateway, key)
+        warm_responses: Final = tuple(
+            _assert_success(_chat(worker, model, key, prompt))
+            for worker, prompt in zip((gateway, peer), prompts[:2], strict=True)
+        )
+        assert len(warm_responses) == 2
+        assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2
+        deleted: Final = gateway.request("POST", "/key/delete", {"keys": [key]})
+        assert deleted.status_code == 200, deleted.text
+        assert _DeleteResponse.model_validate_json(deleted.content).deleted_keys == (key,), deleted.text
+        peer_refusal: Final = eventually(
+            lambda: _chat(peer, model, key, prompts[2]),
+            lambda response: response.status_code == 401,
+            seconds=3,
+        )
+        assert _ErrorResponse.model_validate_json(peer_refusal.content).error.type == "token_not_found_in_db", (
+            peer_refusal.text
+        )
+        wire.drain()
+        fresh_peer_refusal: Final = _chat(peer, model, key, prompts[3])
+        assert fresh_peer_refusal.status_code == 401, fresh_peer_refusal.text
+        assert _ErrorResponse.model_validate_json(fresh_peer_refusal.content).error.type == "token_not_found_in_db", (
+            fresh_peer_refusal.text
+        )
+        fresh_gateway_refusal: Final = _chat(gateway, model, key, prompts[3])
+        assert fresh_gateway_refusal.status_code == 401, fresh_gateway_refusal.text
+        assert _ErrorResponse.model_validate_json(fresh_gateway_refusal.content).error.type == (
+            "token_not_found_in_db"
+        ), fresh_gateway_refusal.text
+        assert wire.drain() == (), "deleted-key requests reached the upstream"

tests/integration/spend/test_key_budget_lockout.py

@@ -1,13 +1,167 @@
+import json
+import os
 import uuid
 from hashlib import sha256
+from pathlib import Path
 from typing import Final
 
 import httpx
 import pytest
 from integration._support.client import Gateway, eventually, object_value
 from integration._support.database import read_rows
+from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, Wire, wire_server
 from openai.types.chat import ChatCompletionChunk
-from pydantic import JsonValue
+from pydantic import BaseModel, JsonValue, TypeAdapter
+from redis import Redis
+
+_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
+_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
+_EXPECTED_MODEL_DISCOVERY: Final = ("GET", "/v1/models")
+_EXPECTED_CHAT: Final = ("POST", "/v1/chat/completions")
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+_FULL_USAGE: Final = _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40)
+_POST_RESET_USAGE: Final = _Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20)
+
+
+class _AssistantMessage(BaseModel):
+    content: str
+
+
+class _Choice(BaseModel):
+    message: _AssistantMessage
+
+
+class _ChatResponse(BaseModel):
+    id: str
+    choices: tuple[_Choice, ...]
+    usage: _Usage
+
+
+class _ModelInfoEntry(BaseModel):
+    model_name: str
+
+
+class _ModelInfoResponse(BaseModel):
+    data: tuple[_ModelInfoEntry, ...]
+
+
+class _BudgetError(BaseModel):
+    type: str
+
+
+class _ErrorResponse(BaseModel):
+    error: _BudgetError
+
+
+class _ResetResponse(BaseModel):
+    key_hash: str
+    spend: float
+    previous_spend: float
+
+
+class _KeyInfo(BaseModel):
+    spend: float
+    max_budget: float
+
+
+class _KeyInfoResponse(BaseModel):
+    info: _KeyInfo
+
+
+class _SpendRow(BaseModel):
+    spend: float
+
+
+class _UserMessage(BaseModel):
+    content: str
+
+
+class _ChatRequest(BaseModel):
+    model: str
+    max_tokens: int
+    messages: tuple[_UserMessage, ...]
+
+
+def _chat_body(model: str, prompt: str, max_tokens: int = 20) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "max_tokens": max_tokens,
+        "messages": [{"role": "user", "content": prompt}],
+    }
+
+
+def _wire_reply(identity: str, usage: _Usage) -> Reply:
+    return Reply(
+        body=json.dumps(
+            {
+                "id": identity,
+                "object": "chat.completion",
+                "created": 1,
+                "model": "integration-keys-auth-reset",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "wire response"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": usage.model_dump(),
+            }
+        ).encode()
+    )
+
+
+def _assert_wire_requests(wire: Wire, expected: tuple[tuple[str, str], ...]) -> None:
+    observed: Final = tuple(
+        sorted(
+            (request.method, request.target)
+            for request in wire.drain()
+            if (request.method, request.target) != _EXPECTED_MODEL_DISCOVERY
+        )
+    )
+    assert observed == tuple(sorted(expected)), observed
+
+
+def _wire_chat(gateway: Gateway, model: str, key: str, prompt: str, max_tokens: int = 20) -> httpx.Response:
+    return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt, max_tokens), key=key)
+
+
+def _assert_wire_success(response: httpx.Response, expected_usage: _Usage) -> _ChatResponse:
+    assert response.status_code == 200, response.text
+    payload: Final = _ChatResponse.model_validate_json(response.content)
+    assert payload.choices[0].message.content == "wire response", response.text
+    assert payload.usage == expected_usage, response.text
+    return payload
+
+
+def _model_is_available(gateway: Gateway, model: str) -> bool:
+    response: Final = gateway.request("GET", "/model/info")
+    if response.status_code != 200:
+        return False
+    models: Final = _ModelInfoResponse.model_validate_json(response.content)
+    return any(entry.model_name == model for entry in models.data)
+
+
+def _spend(gateway: Gateway, key_hash: str) -> float:
+    rows: Final = TypeAdapter(tuple[_SpendRow, ...]).validate_python(
+        read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (key_hash,))
+    )
+    assert len(rows) == 1
+    return rows[0].spend
+
+
+def _assert_budget_refusal(gateway: Gateway, model: str, key: str, prompt: str) -> None:
+    response: Final = _wire_chat(gateway, model, key, prompt)
+    assert response.status_code == 422, response.text
+    assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text
 
 
 def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(gateway: Gateway) -> None:
@@ -113,6 +267,133 @@ def test_raising_a_spent_keys_budget_restores_serving(gateway: Gateway) -> None:
         assert len(upstream.get("/__observations").json()["requests"]) == 3
 
 
+@pytest.mark.parametrize("spelling", ("plaintext", "hashed"))
+def test_reset_spend_lifts_an_exhausted_key_on_both_workers_for_both_path_spellings(
+    gateway: Gateway, tmp_path: Path, spelling: str
+) -> None:
+    backend: Final = "integration-keys-auth-reset"
+    prompts: Final = tuple(f"reset spend case {index} {uuid.uuid4().hex}" for index in range(11))
+    accepted_after_reset: Final = frozenset((prompts[3], prompts[4], prompts[7], prompts[8]))
+
+    def respond(request: Request) -> Reply:
+        if request.method == "GET":
+            assert request.target == "/v1/models", request
+            assert request.headers["authorization"] == "Bearer synthetic-openai-key"
+            assert request.body == b""
+            return Reply(
+                body=json.dumps(
+                    {
+                        "object": "list",
+                        "data": [
+                            {
+                                "id": backend,
+                                "object": "model",
+                                "created": 0,
+                                "owned_by": "integration",
+                            }
+                        ],
+                    }
+                ).encode()
+            )
+        assert request.method == "POST"
+        assert request.target == "/v1/chat/completions"
+        assert request.headers["authorization"] == "Bearer synthetic-openai-key"
+        body: Final = _CHAT_BODY.validate_json(request.body)
+        parsed: Final = _ChatRequest.model_validate_json(request.body)
+        prompt: Final = parsed.messages[0].content
+        max_tokens: Final = 10 if prompt in accepted_after_reset else 20
+        assert parsed.model == backend
+        assert parsed.max_tokens == max_tokens
+        assert len(parsed.messages) == 1 and prompt in prompts, request.body
+        assert body == _chat_body(backend, prompt, max_tokens), body
+        usage: Final = _POST_RESET_USAGE if prompt in accepted_after_reset else _FULL_USAGE
+        return _wire_reply("chatcmpl-" + uuid.uuid4().hex, usage)
+
+    with (
+        wire_server(respond) as wire,
+        gateway.scenario() as scenario,
+        Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache,
+    ):
+        model: Final = scenario.model(
+            model=f"openai/{backend}",
+            api_base=f"{wire.url}/v1",
+            api_key="synthetic-openai-key",
+            input_cost_per_token=0.001,
+            output_cost_per_token=0.002,
+        )
+        subscriber_baseline: Final = cache.pubsub_numsub(_CHANNEL)[0][1]
+        with owned_proxy(gateway, tmp_path / "second-worker", {}) as second:
+            eventually(
+                lambda: cache.pubsub_numsub(_CHANNEL)[0][1],
+                lambda count: count >= subscriber_baseline + 1,
+                seconds=30,
+            )
+            eventually(lambda: _model_is_available(second, model), bool, seconds=30)
+            key: Final = scenario.key(models=[model], max_budget=0.06)
+            key_hash: Final = sha256(key.encode()).hexdigest()
+            first_spend: Final = _assert_wire_success(_wire_chat(gateway, model, key, prompts[0]), _FULL_USAGE)
+            assert first_spend.usage.total_tokens == 40
+            exhausted_spend: Final = eventually(
+                lambda: _spend(gateway, key_hash), lambda spend: spend >= 0.06, seconds=70
+            )
+            assert exhausted_spend == pytest.approx(0.06)
+            _assert_wire_requests(wire, (_EXPECTED_CHAT,))
+            _assert_budget_refusal(gateway, model, key, prompts[1])
+            _assert_budget_refusal(second, model, key, prompts[2])
+            assert wire.drain() == (), "exhausted-key requests reached the upstream"
+
+            identifier: Final = key if spelling == "plaintext" else key_hash
+            for cycle, reset_to in enumerate((0, 0.02)):
+                prior_spend = _spend(gateway, key_hash)
+                reset = gateway.request("POST", f"/key/{identifier}/reset_spend", {"reset_to": reset_to})
+                assert reset.status_code == 200, reset.text
+                reset_result = _ResetResponse.model_validate_json(reset.content)
+                assert reset_result.key_hash == key_hash, reset.text
+                assert reset_result.spend == pytest.approx(reset_to), reset.text
+                assert reset_result.previous_spend == pytest.approx(prior_spend), reset.text
+                info_after_reset = gateway.request("GET", "/key/info", params={"key": key})
+                assert info_after_reset.status_code == 200, info_after_reset.text
+                assert _KeyInfoResponse.model_validate_json(info_after_reset.content).info.spend == pytest.approx(
+                    reset_to
+                ), info_after_reset.text
+                assert _spend(gateway, key_hash) == pytest.approx(reset_to)
+                peer_info_after_reset = eventually(
+                    lambda: second.request("GET", "/key/info", params={"key": key}),
+                    lambda response, reset_to=reset_to: (
+                        response.status_code == 200
+                        and _KeyInfoResponse.model_validate_json(response.content).info.spend == pytest.approx(reset_to)
+                    ),
+                    seconds=10,
+                )
+                assert _KeyInfoResponse.model_validate_json(peer_info_after_reset.content).info.spend == pytest.approx(
+                    reset_to
+                ), peer_info_after_reset.text
+                workers = (gateway, second) if spelling == "plaintext" else (second, gateway)
+                first_chat = _wire_chat(workers[0], model, key, prompts[3 + cycle * 4], max_tokens=10)
+                _assert_wire_success(first_chat, _POST_RESET_USAGE)
+                first_reported_spend = float(first_chat.headers["x-litellm-key-spend"])
+                assert first_reported_spend == pytest.approx(reset_to + 0.03), first_chat.headers
+                eventually(
+                    lambda: float(cache.get(f"spend:key:{key_hash}") or 0),
+                    lambda counter, reset_to=reset_to: counter == pytest.approx(reset_to + 0.03),
+                    seconds=30,
+                )
+                second_chat = _wire_chat(workers[1], model, key, prompts[4 + cycle * 4], max_tokens=10)
+                _assert_wire_success(second_chat, _POST_RESET_USAGE)
+                second_reported_spend = float(second_chat.headers["x-litellm-key-spend"])
+                assert reset_to + 0.03 <= second_reported_spend <= reset_to + 0.06, second_chat.headers
+                re_exhausted_spend = eventually(
+                    lambda: _spend(gateway, key_hash),
+                    lambda spend, reset_to=reset_to: spend >= reset_to + 0.06,
+                    seconds=70,
+                )
+                assert re_exhausted_spend == pytest.approx(reset_to + 0.06)
+                _assert_wire_requests(wire, (_EXPECTED_CHAT, _EXPECTED_CHAT))
+                _assert_budget_refusal(gateway, model, key, prompts[5 + cycle * 4])
+                _assert_budget_refusal(second, model, key, prompts[6 + cycle * 4])
+                assert wire.drain() == (), "re-exhausted-key requests reached the upstream"
+
+
 def _denied_with_budget_exceeded(gateway: Gateway, model: str, key: str) -> None:
     denied: Final = _bounded_chat(gateway, model, key)
     assert denied.status_code == 422, denied.text