diff --git a/tests/e2e/access_control/test_access_control_e2e.py b/tests/e2e/access_control/test_access_control_e2e.py index af7e9a099fd..9d01f2915e7 100644 --- a/tests/e2e/access_control/test_access_control_e2e.py +++ b/tests/e2e/access_control/test_access_control_e2e.py @@ -24,7 +24,7 @@ from access_control_client import ( from e2e_config import unique_marker from e2e_http import Success, UnauthorizedError, UnknownApiError, unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, LiteLLMParamsBody from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -32,6 +32,7 @@ pytestmark = pytest.mark.e2e ALLOWED_MODEL = "gemini-2.5-flash" DISALLOWED_MODEL = "gpt-5.5" VIRTUAL_KEY_BACKEND = "anthropic/claude-haiku-4-5-20251001" +EMBEDDING_MODEL = "openai-text-embedding-3-small" class TestAccessControl: @@ -71,6 +72,31 @@ class TestAccessControl: f"403 body must be a model-access denial, got: {result.body[:300]}" ) + @pytest.mark.covers("other.auth.virtual_key.route_group_allowed") + def test_llm_api_routes_group_grants_every_llm_endpoint( + self, client: AccessControlClient, resources: ResourceManager + ) -> None: + key = client.llm_only_key() + resources.defer(lambda: client.delete_key(key)) + + chat = client.chat_status(key, ALLOWED_MODEL, f"capital of France? {unique_marker()}") + assert chat.status_code == 200, ( + f"llm_api_routes key must reach /chat/completions, got {chat.status_code}: {chat.body[:300]}" + ) + assert ChatResponse.model_validate_json(chat.body).choices, ( + f"200 must carry a real completion, not an error envelope: {chat.body[:300]}" + ) + + embedding = unwrap( + client.proxy.embed(key, EmbedBody(model=EMBEDDING_MODEL, input=f"route group {unique_marker()}")) + ) + assert embedding.model, f"llm_api_routes key reached /embeddings but got no model back: {embedding}" + + denied = client.create_model_status(key, f"e2e-route-group-{unique_marker()}") + assert denied.status_code == 403 and ROUTE_NOT_ALLOWED_MARKER in denied.body, ( + f"the same key must still be shut out of /model/new, got {denied.status_code}: {denied.body[:300]}" + ) + def test_llm_only_key_forbidden_from_management_route_403( self, client: AccessControlClient, resources: ResourceManager ) -> None: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 387280c8023..2b897f5f07f 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -40,6 +40,8 @@ from models import ( KeyListParams, KeyListResponse, KeyRegenerateBody, + KeyResetSpendBody, + KeyResetSpendResponse, KeyUpdateBody, ModelDeleteBody, OrgDeleteBody, @@ -191,16 +193,26 @@ class ManagementClient: response_type=NoBody, ) ) - def regenerate_key(self, key: str) -> str: + def regenerate_key(self, key: str, *, grace_period: str | None = None) -> str: return unwrap( self.proxy.transport.post( "/key/regenerate", headers=self.proxy.transport.master, - json=KeyRegenerateBody(key=key), + json=KeyRegenerateBody(key=key, grace_period=grace_period), response_type=KeyGenerateResponse, ) ).key + def reset_key_spend(self, key: str, reset_to: float) -> KeyResetSpendResponse: + return unwrap( + self.proxy.transport.post( + f"/key/{key}/reset_spend", + headers=self.proxy.transport.master, + json=KeyResetSpendBody(reset_to=reset_to), + response_type=KeyResetSpendResponse, + ) + ) + def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]: """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is who is asking: the master key by default, or a virtual key.""" diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index 711175abb0d..8b7d5f0eb6f 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -18,7 +18,7 @@ from typing import Literal import pytest from e2e_config import unique_marker -from e2e_http import NoBody, unwrap +from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient from models import KeyDeleteBody, KeyGenerateBody, KeyUpdateBody @@ -26,6 +26,9 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e +TINY_BUDGET = 3e-6 +SPEND_MODEL = "claude-haiku-4-5" + class KeyToggleBlockBody(BaseModel): key: str @@ -82,6 +85,30 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke return key +def _is_budget_block(outcome: StreamingResponse) -> bool: + return not outcome.ok and "budget_exceeded" in outcome.body + + +def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None: + for _ in range(40): + outcome = client.chat_status(key, SPEND_MODEL, f"spend {unique_marker()}") + if _is_budget_block(outcome): + assert outcome.status_code == 429, ( + f"budget refusal must be 429, got {outcome.status_code}: {outcome.body[:200]}" + ) + return + assert outcome.ok, f"paid call failed before the budget tripped ({outcome.status_code}): {outcome.body[:300]}" + time.sleep(2) + pytest.fail(f"max_budget={TINY_BUDGET} never blocked a call on the key") + + +def _settled_spend(client: ManagementClient, key: str) -> float | None: + first = client.proxy.key_info(key).spend or 0.0 + time.sleep(client.proxy.poll_interval) + second = client.proxy.key_info(key).spend or 0.0 + return second if first > 0 and first == second else None + + def _block(client: ManagementClient, key: str) -> None: _ = unwrap( client.proxy.transport.post( @@ -197,6 +224,32 @@ class TestKeyManagementRoutes: "/key/info never reported max_budget 42.0 after /key/bulk_update before the deadline", ) + @pytest.mark.covers("other.key_mgmt.spend_reset.resets_to_value") + def test_reset_spend_zeroes_recorded_spend_and_lifts_the_budget_block( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + key = _generate_key(client, resources, KeyGenerateBody(models=[SPEND_MODEL], max_budget=TINY_BUDGET)) + _spend_until_budget_blocks(client, key) + recorded = _poll( + client, lambda: _settled_spend(client, key), "key spend never landed in /key/info before the deadline" + ) + + reset = client.reset_key_spend(key, reset_to=0.0) + assert reset.previous_spend == recorded, ( + f"reset_spend reported previous_spend {reset.previous_spend}, /key/info had recorded {recorded}" + ) + assert reset.spend == 0.0, f"reset_spend to 0 reported spend {reset.spend}" + assert client.proxy.key_info(key).spend == 0.0, "/key/info still reports spend after the reset to 0" + + def call_allowed_again() -> bool | None: + outcome = client.chat_status(key, SPEND_MODEL, f"after reset {unique_marker()}") + if _is_budget_block(outcome): + return None + assert outcome.ok, f"post-reset call failed ({outcome.status_code}): {outcome.body[:300]}" + return True + + _ = _poll(client, call_allowed_again, "the key stayed budget-blocked after its spend was reset to 0") + @pytest.mark.covers("mgmt.key.generate.admin_only") def test_generate_forbidden_for_non_admin_key( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index a56eb853823..476165b715d 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -12,6 +12,7 @@ from __future__ import annotations import math import time from collections.abc import Callable +from typing import Final import pytest @@ -42,6 +43,10 @@ from models import ( pytestmark = pytest.mark.e2e +REGENERATE_GRACE_PERIOD = "15s" +REGENERATE_GRACE_SECONDS = 15.0 + + def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: @@ -365,6 +370,36 @@ class TestKeyRegeneration: client, old_rejected, "old key was still accepted after regeneration (never rejected 401) at the deadline" ) + @pytest.mark.covers("other.key_mgmt.regenerate.grace_period_honored") + def test_regenerate_with_grace_period_keeps_old_key_until_revoked( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + old_key = _generate_key(client, resources, KeyGenerateBody(models=["gpt-5.5"])) + + new_key = client.regenerate_key(old_key, grace_period=REGENERATE_GRACE_PERIOD) + resources.defer(lambda: client.proxy.delete_key(new_key)) + revoke_at: Final = time.monotonic() + REGENERATE_GRACE_SECONDS + assert new_key != old_key, "regenerate returned the same key string, so no rotation happened" + + def old_accepted() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.ok else None + + _ = _poll(client, old_accepted, "old key was rejected 401 inside its grace period at the deadline") + assert time.monotonic() < revoke_at, ( + f"old key was only accepted after its {REGENERATE_GRACE_PERIOD} grace period had elapsed" + ) + + def old_rejected() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code == 401 else None + + _ = _poll( + client, + old_rejected, + f"old key was still accepted past its {REGENERATE_GRACE_PERIOD} grace period (never 401) at the deadline", + ) + class TestTeamRoutes: @pytest.mark.covers("mgmt.team.new.persists") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b5229744d6f..5de49ead3ed 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -83,6 +83,16 @@ class KeyGenerateResponse(BaseModel): class KeyRegenerateBody(BaseModel): key: str + grace_period: str | None = None + + +class KeyResetSpendBody(BaseModel): + reset_to: float + + +class KeyResetSpendResponse(BaseModel): + spend: float + previous_spend: float class KeyDeleteBody(BaseModel):