diff --git a/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py b/tests/integration/authorization/test_jwt_key_mapping_lifecycle.py new file mode 100644 index 00000000000..03ca94dd561 --- /dev/null +++ b/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), + ) diff --git a/tests/integration/authorization/test_warmed_policy.py b/tests/integration/authorization/test_warmed_policy.py index f114e78507a..b4ffc10c3c4 100644 --- a/tests/integration/authorization/test_warmed_policy.py +++ b/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" diff --git a/tests/integration/database/test_reader_writer_regeneration.py b/tests/integration/database/test_reader_writer_regeneration.py index 4161d0b04a5..62c152167ab 100644 --- a/tests/integration/database/test_reader_writer_regeneration.py +++ b/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,) diff --git a/tests/integration/management/test_key_delete_revocation.py b/tests/integration/management/test_key_delete_revocation.py new file mode 100644 index 00000000000..5b3527a52ae --- /dev/null +++ b/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" diff --git a/tests/integration/spend/test_key_budget_lockout.py b/tests/integration/spend/test_key_budget_lockout.py index 7108978f6e5..56b4db89bcb 100644 --- a/tests/integration/spend/test_key_budget_lockout.py +++ b/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