diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 20ab0212854..2d7e3f5c885 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -13,7 +13,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `realtime/` - realtime websocket sessions, including the pipecat audio path - `budgets/` - budget definition, enforcement, and reset windows (key, team, tag, soft, multi-window) - `spend_tracking/` - spend logging and cost attribution on `/spend/*` -- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials +- `management/` - key/team/user/organization and model-deployment management routes: create/update/delete persistence via the info routes, team membership, tpm persistence, and llm-only-key route denials - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (rate limits, fallbacks, cooldowns) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 5520b44993d..b5a6794363c 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -1,7 +1,7 @@ """Client for the management-routes e2e suite: the shared Gateway plus the -key/team/user/organization writes, the info/list read-backs the tests assert, -and the raw-status calls judged by HTTP outcome (chat under a scoped key, an -llm-only key hitting a management route). +key/team/user/organization/model-deployment writes, the info/list read-backs +the tests assert, and the raw-status calls judged by HTTP outcome (chat under a +scoped key, an llm-only key hitting a management route). """ from __future__ import annotations @@ -18,6 +18,13 @@ from models import ( KeyListParams, KeyListResponse, KeyUpdateBody, + LiteLLMParamsBody, + ModelDeleteBody, + ModelInfoBody, + ModelInfoEntry, + ModelNewBody, + ModelUpdateBody, + ModelUpdateParams, OrgDeleteBody, OrgInfoParams, OrgInfoResponse, @@ -43,6 +50,18 @@ from models import ( MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" +UNKNOWN_MODEL_MARKER = "Invalid model name passed in model=" +NO_DEPLOYMENTS_MARKER = "There are no healthy deployments" + + +def is_deleted_model_rejection(outcome: StreamingResponse, model: str) -> bool: + """True if the gateway refused the call because `model` is gone: a 400 naming + the model, either the proxy's unknown-model shape (the data plane never knew + the group) or the router's no-healthy-deployments shape (the group name + outlives its last deployment in the router until restart).""" + if outcome.status_code != 400 or model not in outcome.body: + return False + return UNKNOWN_MODEL_MARKER in outcome.body or NO_DEPLOYMENTS_MARKER in outcome.body @dataclass(frozen=True, slots=True) @@ -84,6 +103,50 @@ class ManagementClient: ) ).total_count + def find_deployment(self, model_name: str) -> ModelInfoEntry | None: + return next( + (entry for entry in self.gateway.model_info() if entry.model_name == model_name), + None, + ) + + def update_model_tpm(self, model_id: str, tpm: int) -> None: + _ = unwrap( + self.gateway.transport.post( + "/model/update", + headers=self.gateway.transport.master, + json=ModelUpdateBody( + litellm_params=ModelUpdateParams(tpm=tpm), + model_info=ModelInfoBody(id=model_id), + ), + response_type=NoBody, + ) + ) + + def delete_model_strict(self, model_id: str) -> None: + """Strict delete for the act phase of a test: a failed delete is a hard + failure, unlike the warn-only Gateway.delete_model used at teardown.""" + _ = unwrap( + self.gateway.transport.post( + "/model/delete", + headers=self.gateway.transport.master, + json=ModelDeleteBody(id=model_id), + response_type=NoBody, + ) + ) + + def create_model_status( + self, key: str, model_name: str, litellm_params: LiteLLMParamsBody + ) -> StreamingResponse: + return self.gateway.transport.send( + "/model/new", + headers=self.gateway.transport.bearer(key), + json=ModelNewBody( + model_name=model_name, + litellm_params=litellm_params, + model_info=ModelInfoBody(id=model_name), + ), + ) + def create_team(self, body: TeamNewBody) -> str: return unwrap( self.gateway.transport.post( diff --git a/tests/e2e/management/test_model_management_e2e.py b/tests/e2e/management/test_model_management_e2e.py new file mode 100644 index 00000000000..e5ba35cf4e2 --- /dev/null +++ b/tests/e2e/management/test_model_management_e2e.py @@ -0,0 +1,183 @@ +"""Live e2e: the model-management routes' add / update / delete contract. + +Each test provisions its own deployment through /model/new under a unique name +(deleted on teardown) and asserts both halves of the lifecycle contract: the +recorded state (/model/info reflects the write) and the enforced behavior (the +gateway serves or refuses traffic accordingly). Management writes land on the +control plane while chat is served by the data plane, which picks the change up +on its DB sync, so the traffic-facing read-backs poll to a deadline instead of +asserting once. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable + +import pytest + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, Success +from lifecycle import ResourceManager +from management_client import ( + ROUTE_NOT_ALLOWED_MARKER, + ManagementClient, + is_deleted_model_rejection, +) +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, ModelInfoEntry + +pytestmark = pytest.mark.e2e + + +def _chat_body(model: str) -> ChatBody: + return ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"reply with one word {unique_marker()}")], + max_tokens=128, + ) + + +def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: + deadline = time.monotonic() + client.gateway.poll_timeout + while time.monotonic() < deadline: + found = attempt() + if found is not None: + return found + time.sleep(client.gateway.poll_interval) + pytest.fail(failure) + + +def _delete_if_present(client: ManagementClient, model_name: str) -> None: + if client.find_deployment(model_name) is not None: + client.gateway.delete_model(model_name) + + +def _provision( + client: ManagementClient, resources: ResourceManager, *, tpm: int | None = None +) -> str: + """Register a fresh gemini-backed deployment (id == name, cleaned up on + teardown) and return its unique model name.""" + model_name = f"e2e-mgmt-{unique_marker()}" + _ = client.gateway.create_model( + model_name, + LiteLLMParamsBody(model="gemini/gemini-3.5-flash", api_key="os.environ/GEMINI_API_KEY", tpm=tpm), + ) + resources.defer(lambda: _delete_if_present(client, model_name)) + return model_name + + +def _poll_chat_ok(client: ManagementClient, key: str, model: str) -> ChatResponse: + def attempt() -> ChatResponse | None: + match client.gateway.chat(key, _chat_body(model)): + case Success(data=data): + return data + case _: + return None + + return _poll( + client, + attempt, + f"{model} never became callable through the data plane before the deadline", + ) + + +def _poll_chat_rejects_deleted_model(client: ManagementClient, key: str, model: str) -> None: + def attempt() -> StreamingResponse | None: + outcome = client.chat_status(key, model, f"say hi {unique_marker()}") + return outcome if is_deleted_model_rejection(outcome, model) else None + + _ = _poll( + client, + attempt, + f"deleted model {model} was still served (never rejected with a 400 naming it) at the deadline", + ) + + +class TestModelManagementRoutes: + @pytest.mark.covers("mgmt.model.add.persists") + def test_add_persists_to_model_info_and_serves_chat( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = _provision(client, resources, tpm=500000) + + entry = _poll( + client, + lambda: client.find_deployment(model), + f"{model} never appeared in /model/info after /model/new", + ) + assert entry.litellm_params.model == "gemini/gemini-3.5-flash", ( + f"/model/info reports backend {entry.litellm_params.model!r}, " + f"configured 'gemini/gemini-3.5-flash'" + ) + assert entry.litellm_params.tpm == 500000, ( + f"/model/info reports tpm {entry.litellm_params.tpm}, configured 500000" + ) + + response = _poll_chat_ok(client, scoped_key, model) + assert response.choices, f"chat on {model} succeeded but returned no choices: {response}" + + @pytest.mark.covers("mgmt.model.update.persists") + def test_update_tpm_persists_and_deployment_still_serves( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = _provision(client, resources, tpm=500000) + + client.update_model_tpm(model, 424242) + + def updated() -> ModelInfoEntry | None: + entry = client.find_deployment(model) + if entry is not None and entry.litellm_params.tpm == 424242: + return entry + return None + + entry = _poll( + client, + updated, + f"/model/info never reflected tpm 424242 for {model} after /model/update", + ) + assert entry.litellm_params.model == "gemini/gemini-3.5-flash", ( + f"/model/update merge lost the backend model: {entry.litellm_params.model!r}" + ) + + response = _poll_chat_ok(client, scoped_key, model) + assert response.choices, ( + f"updated deployment {model} no longer serves chat: {response}" + ) + + @pytest.mark.covers("mgmt.model.delete.persists") + def test_delete_removes_from_model_info_and_rejects_chat( + self, client: ManagementClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = _provision(client, resources) + _ = _poll_chat_ok(client, scoped_key, model) + + client.delete_model_strict(model) + + def absent() -> bool | None: + return True if client.find_deployment(model) is None else None + + _ = _poll(client, absent, f"{model} still listed in /model/info after /model/delete") + _poll_chat_rejects_deleted_model(client, scoped_key, model) + + @pytest.mark.covers("mgmt.model.add.member_forbidden") + def test_add_forbidden_for_llm_only_key( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + key = client.llm_only_key() + resources.defer(lambda: client.gateway.delete_key(key)) + model = f"e2e-mgmt-forbidden-{unique_marker()}" + + outcome = client.create_model_status( + key, model, LiteLLMParamsBody(model="gemini/gemini-3.5-flash") + ) + + assert outcome.status_code == 403, ( + f"llm-only key POSTing /model/new must be denied 403, got " + f"{outcome.status_code}: {outcome.body[:300]}" + ) + assert ROUTE_NOT_ALLOWED_MARKER in outcome.body, ( + f"403 body must be a route-permission denial, got: {outcome.body[:300]}" + ) + assert client.find_deployment(model) is None, ( + f"{model} was created despite the 403 route denial" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 0490db286ea..8c542f04567 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -321,6 +321,70 @@ class CustomPricing(BaseModel): return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token +class DeploymentParams(BaseModel): + """The configured litellm_params a /model/info row reports for a deployment, + mirroring litellm's LiteLLM_Params (litellm/types/router.py) field for field + so tests can assert any configured knob. Same pricing field names as + CustomPricing, so pricing tests read them unchanged. Secrets (api_key et al) + come back encrypted, so tests assert presence, never the value. Deliberately + omitted because they have no typed JSON shape to pin: mock_response, + model_info (surfaced as ModelInfoEntry.model_info), the *_router_config + blobs, and configurable_clientside_auth_params.""" + + model_config = ConfigDict(extra="ignore", protected_namespaces=()) + + model: str | None = None + custom_llm_provider: str | None = None + + api_key: str | None = None + api_base: str | None = None + api_version: str | None = None + organization: str | None = None + litellm_credential_name: str | None = None + + tpm: int | None = None + rpm: int | None = None + itpm: int | None = None + otpm: int | None = None + max_parallel_requests: int | None = None + order: int | None = None + weight: int | None = None + + timeout: float | str | None = None + stream_timeout: float | str | None = None + max_retries: int | None = None + + max_budget: float | None = None + budget_duration: str | None = None + default_api_key_tpm_limit: int | None = None + default_api_key_rpm_limit: int | None = None + + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + input_cost_per_second: float | None = None + output_cost_per_second: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + + tags: list[str] | None = None + tag_regex: list[str] | None = None + + use_in_pass_through: bool | None = None + use_litellm_proxy: bool | None = None + use_chat_completions_api: bool | None = None + use_xai_oauth: bool | None = None + merge_reasoning_content_in_choices: bool | None = None + + region_name: str | None = None + aws_region_name: str | None = None + vertex_project: str | None = None + vertex_location: str | None = None + watsonx_region_name: str | None = None + + max_file_size_mb: float | None = None + litellm_trace_id: str | None = None + + class ModelInfoEntry(BaseModel): """One /model/info row. `litellm_params` is the configured deployment (carries any custom-pricing override); `model_info` is the price the proxy resolved for @@ -328,7 +392,7 @@ class ModelInfoEntry(BaseModel): model_config = ConfigDict(protected_namespaces=()) model_name: str - litellm_params: CustomPricing = CustomPricing() + litellm_params: DeploymentParams = DeploymentParams() model_info: CustomPricing = CustomPricing() @@ -388,6 +452,7 @@ class LiteLLMParamsBody(BaseModel): aws_batch_role_arn: str | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None + tpm: int | None = None ModelMode = Literal["batch", "realtime", "image_generation"] @@ -414,6 +479,19 @@ class ModelDeleteBody(BaseModel): id: str +class ModelUpdateParams(BaseModel): + """POST /model/update litellm_params: only the fields being changed; the proxy + merges them over the deployment's stored params.""" + + tpm: int | None = None + + +class ModelUpdateBody(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + litellm_params: ModelUpdateParams + model_info: ModelInfoBody + + # ---------- key / team / user / organization management ----------