mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
8741e8a1ad
The test claiming mgmt.key.generate.happy_path signed in and then only read /key/list, so nothing proved the session key an admin's sign-in mints is actually accepted on /key/generate. It now does what an admin filling in Create New Key does: POST /key/generate under the session key, read the new key back from /key/info, see it in the dashboard's own /key/list, and drive real traffic through it to confirm its model scope is enforced. Adds ManagementClient.generate_key with the same caller_key seam update_key and key_list already use, so the suite can call the route as the master key or as a virtual key. Also wraps the over-long models import. Refusing the dashboard session key on /key/generate turns only this test red; the master-key generate, the key edit, and regenerate stay green.
546 lines
19 KiB
Python
546 lines
19 KiB
Python
"""Client for the management-routes e2e suite: the shared ProxyClient 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).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
import jwt
|
|
|
|
from e2e_config import MASTER_KEY
|
|
from proxy_client import ProxyClient
|
|
from e2e_http import (
|
|
AuthHeaders,
|
|
NetworkError,
|
|
NoBody,
|
|
ProbeResult,
|
|
Result,
|
|
StreamingResponse,
|
|
Success,
|
|
UnknownApiError,
|
|
unwrap,
|
|
)
|
|
from models import (
|
|
ChatBody,
|
|
ChatMessage,
|
|
ConnectionTestBody,
|
|
ConnectionTestResponse,
|
|
CustomerDeleteBody,
|
|
CustomerInfoParams,
|
|
CustomerNewBody,
|
|
CustomerResponse,
|
|
KeyBlockBody,
|
|
KeyDeleteBody,
|
|
KeyGenerateBody,
|
|
KeyGenerateResponse,
|
|
KeyListParams,
|
|
KeyListResponse,
|
|
KeyRegenerateBody,
|
|
KeyUpdateBody,
|
|
ModelDeleteBody,
|
|
OrgDeleteBody,
|
|
OrgInfoParams,
|
|
OrgInfoResponse,
|
|
OrgNewBody,
|
|
OrgNewResponse,
|
|
OrgUpdateBody,
|
|
TagDeleteBody,
|
|
TagListEntry,
|
|
TagListResponse,
|
|
TagNewBody,
|
|
TeamData,
|
|
TeamDeleteBody,
|
|
TeamInfoParams,
|
|
TeamInfoResponse,
|
|
TeamListResponse,
|
|
TeamMemberAddBody,
|
|
TeamMemberDeleteBody,
|
|
TeamMemberEntry,
|
|
TeamNewBody,
|
|
TeamNewResponse,
|
|
TeamUpdateBody,
|
|
UiLoginBody,
|
|
UiLoginResponse,
|
|
UiSessionClaims,
|
|
UserDeleteBody,
|
|
UserDeleteResponse,
|
|
UserInfoParams,
|
|
UserInfoResponse,
|
|
UserListParams,
|
|
UserListResponse,
|
|
UserNewBody,
|
|
UserNewResponse,
|
|
UserUpdateBody,
|
|
)
|
|
|
|
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
|
|
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
|
DASHBOARD_SESSION_TEAM_ID = "litellm-dashboard"
|
|
_TEAM_READY_ATTEMPTS = 15
|
|
_TEAM_READY_SLEEP_SECONDS = 0.4
|
|
_KEY_WRITE_ATTEMPTS = 5
|
|
_TRANSIENT_BACKEND_MARKERS = ("connecting to redis", "name resolution")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DashboardSession:
|
|
"""What a dashboard sign-in hands the Admin UI: the session key it sends as
|
|
its bearer on every subsequent call, the claims it renders the signed-in user
|
|
from, and where it lands the browser."""
|
|
|
|
session_key: str
|
|
claims: UiSessionClaims
|
|
redirect_url: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ManagementClient:
|
|
proxy: ProxyClient
|
|
master_key: str
|
|
|
|
def llm_only_key(self) -> str:
|
|
return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
|
|
|
|
def generate_key(self, body: KeyGenerateBody, *, caller_key: str | None = None) -> Result[KeyGenerateResponse]:
|
|
"""POST /key/generate. `caller_key` is who is creating the key: the master
|
|
key by default, or a virtual key (an admin filling in Create New Key on the
|
|
dashboard creates it under the session key their sign-in minted). Returns
|
|
the outcome rather than unwrapping it, so a caller can poll a route that is
|
|
only transiently refusing."""
|
|
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
|
return self.proxy.transport.post(
|
|
"/key/generate",
|
|
headers=headers,
|
|
json=body,
|
|
response_type=KeyGenerateResponse,
|
|
)
|
|
|
|
def update_key(self, body: KeyUpdateBody, *, caller_key: str | None = None) -> Result[NoBody]:
|
|
"""POST /key/update. `caller_key` is who is editing: the master key by
|
|
default, or a virtual key (the dashboard edits under the session key its
|
|
sign-in minted, never the master key). Returns the outcome rather than
|
|
unwrapping it, so a caller can poll a route that is only transiently
|
|
refusing; `update_key_models` is the unwrapping shorthand."""
|
|
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
|
last: Result[NoBody] = NetworkError(message="/key/update was never attempted")
|
|
for attempt in range(_KEY_WRITE_ATTEMPTS):
|
|
last = self.proxy.transport.post(
|
|
"/key/update",
|
|
headers=headers,
|
|
json=body,
|
|
response_type=NoBody,
|
|
)
|
|
match last:
|
|
case UnknownApiError(body=error_body) if any(
|
|
marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS
|
|
):
|
|
time.sleep(0.5 * (attempt + 1))
|
|
continue
|
|
case _:
|
|
break
|
|
return last
|
|
|
|
def update_key_models(self, key: str, models: list[str]) -> None:
|
|
_ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models)))
|
|
|
|
def delete_key_strict(self, key: str) -> None:
|
|
"""Strict delete for the act phase of a test: a failed delete is a hard
|
|
failure, unlike the warn-only ProxyClient.delete_key used at teardown."""
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/key/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=KeyDeleteBody(keys=[key]),
|
|
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 ProxyClient.delete_model used at teardown."""
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/model/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=ModelDeleteBody(id=model_id),
|
|
response_type=NoBody,
|
|
)
|
|
)
|
|
|
|
def connection_test(self, body: ConnectionTestBody) -> Result[ConnectionTestResponse]:
|
|
"""POST /health/test_connection, the call behind the Admin UI's Test
|
|
Connection button, probing the live provider with the supplied params."""
|
|
return self.proxy.transport.post(
|
|
"/health/test_connection",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=ConnectionTestResponse,
|
|
timeout=120.0,
|
|
)
|
|
|
|
def block_key(self, key: str) -> None:
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/key/block",
|
|
headers=self.proxy.transport.master,
|
|
json=KeyBlockBody(key=key),
|
|
response_type=NoBody,
|
|
)
|
|
)
|
|
def regenerate_key(self, key: str) -> str:
|
|
return unwrap(
|
|
self.proxy.transport.post(
|
|
"/key/regenerate",
|
|
headers=self.proxy.transport.master,
|
|
json=KeyRegenerateBody(key=key),
|
|
response_type=KeyGenerateResponse,
|
|
)
|
|
).key
|
|
|
|
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."""
|
|
headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key)
|
|
return self.proxy.transport.get(
|
|
"/key/list",
|
|
headers=headers,
|
|
params=KeyListParams(key_alias=key_alias),
|
|
response_type=KeyListResponse,
|
|
)
|
|
|
|
def key_alias_count(self, key_alias: str) -> int:
|
|
return unwrap(self.key_list(key_alias)).total_count
|
|
|
|
def dashboard_login(self, username: str, password: str) -> DashboardSession:
|
|
"""POST /v2/login, the call the Admin UI's sign-in form makes.
|
|
|
|
The proxy authenticates the credentials, mints a UI session key for the
|
|
signed-in user, and hands it back inside a JWT signed with the master key.
|
|
Decoding that JWT is the only way to reach the session key, and it is what
|
|
the dashboard itself does before it can call a single management route."""
|
|
response = unwrap(
|
|
self.proxy.transport.post(
|
|
"/v2/login",
|
|
headers=AuthHeaders(),
|
|
json=UiLoginBody(username=username, password=password),
|
|
response_type=UiLoginResponse,
|
|
)
|
|
)
|
|
decoded: object = jwt.decode(response.token, self.master_key, algorithms=["HS256"])
|
|
claims = UiSessionClaims.model_validate(decoded)
|
|
return DashboardSession(
|
|
session_key=claims.key,
|
|
claims=claims,
|
|
redirect_url=response.redirect_url,
|
|
)
|
|
|
|
def create_team(self, body: TeamNewBody) -> str:
|
|
team_id = unwrap(
|
|
self.proxy.transport.post(
|
|
"/team/new",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=TeamNewResponse,
|
|
)
|
|
).team_id
|
|
self._wait_for_team(team_id)
|
|
return team_id
|
|
|
|
def update_team(self, body: TeamUpdateBody) -> None:
|
|
last: Result[NoBody] | None = None
|
|
for attempt in range(5):
|
|
last = self.proxy.transport.post(
|
|
"/team/update",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=NoBody,
|
|
)
|
|
match last:
|
|
case Success():
|
|
return
|
|
case UnknownApiError(body=body_text) if (
|
|
"connecting to redis" in body_text.lower() or "name resolution" in body_text.lower()
|
|
):
|
|
time.sleep(0.5 * (attempt + 1))
|
|
continue
|
|
case _:
|
|
break
|
|
assert last is not None
|
|
raise AssertionError(last)
|
|
|
|
def delete_team(self, team_id: str) -> None:
|
|
_ = self.proxy.transport.post(
|
|
"/team/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=TeamDeleteBody(team_ids=[team_id]),
|
|
response_type=NoBody,
|
|
)
|
|
|
|
def team_info(self, team_id: str) -> TeamData:
|
|
return unwrap(
|
|
self.proxy.transport.get(
|
|
"/team/info",
|
|
headers=self.proxy.transport.master,
|
|
params=TeamInfoParams(team_id=team_id),
|
|
response_type=TeamInfoResponse,
|
|
)
|
|
).team_info
|
|
|
|
def team_list_ids(self) -> tuple[str, ...]:
|
|
return tuple(
|
|
entry.team_id
|
|
for entry in unwrap(
|
|
self.proxy.transport.get(
|
|
"/team/list",
|
|
headers=self.proxy.transport.master,
|
|
params=NoBody(),
|
|
response_type=TeamListResponse,
|
|
)
|
|
).root
|
|
)
|
|
|
|
def team_info_status(self, team_id: str) -> ProbeResult:
|
|
return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id))
|
|
|
|
def _wait_for_team(self, team_id: str) -> None:
|
|
last: Result[TeamInfoResponse] | None = None
|
|
for _ in range(_TEAM_READY_ATTEMPTS):
|
|
last = self.proxy.transport.get(
|
|
"/team/info",
|
|
headers=self.proxy.transport.master,
|
|
params=TeamInfoParams(team_id=team_id),
|
|
response_type=TeamInfoResponse,
|
|
)
|
|
match last:
|
|
case Success():
|
|
return
|
|
case _:
|
|
time.sleep(_TEAM_READY_SLEEP_SECONDS)
|
|
assert last is not None
|
|
raise AssertionError(last)
|
|
|
|
def add_team_member(self, team_id: str, user_id: str) -> None:
|
|
last: Result[NoBody] | None = None
|
|
for attempt in range(_TEAM_READY_ATTEMPTS):
|
|
last = self.proxy.transport.post(
|
|
"/team/member_add",
|
|
headers=self.proxy.transport.master,
|
|
json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)),
|
|
response_type=NoBody,
|
|
)
|
|
match last:
|
|
case Success():
|
|
return
|
|
case UnknownApiError(body=body) if (
|
|
"doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS
|
|
):
|
|
time.sleep(_TEAM_READY_SLEEP_SECONDS)
|
|
continue
|
|
case _:
|
|
break
|
|
assert last is not None
|
|
raise AssertionError(last)
|
|
|
|
def delete_team_member(self, team_id: str, user_id: str) -> None:
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/team/member_delete",
|
|
headers=self.proxy.transport.master,
|
|
json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id),
|
|
response_type=NoBody,
|
|
)
|
|
)
|
|
|
|
def create_user(self, body: UserNewBody) -> str:
|
|
return unwrap(
|
|
self.proxy.transport.post(
|
|
"/user/new",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=UserNewResponse,
|
|
)
|
|
).user_id
|
|
|
|
def create_customer(self, user_id: str) -> str:
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/customer/new",
|
|
headers=self.proxy.transport.master,
|
|
json=CustomerNewBody(user_id=user_id),
|
|
response_type=CustomerResponse,
|
|
)
|
|
)
|
|
return user_id
|
|
|
|
def customer_info(self, end_user_id: str) -> CustomerResponse:
|
|
return unwrap(
|
|
self.proxy.transport.get(
|
|
"/customer/info",
|
|
headers=self.proxy.transport.master,
|
|
params=CustomerInfoParams(end_user_id=end_user_id),
|
|
response_type=CustomerResponse,
|
|
)
|
|
)
|
|
|
|
def delete_customer(self, user_id: str) -> None:
|
|
_ = self.proxy.transport.post(
|
|
"/customer/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=CustomerDeleteBody(user_ids=[user_id]),
|
|
response_type=NoBody,
|
|
)
|
|
|
|
def update_user(self, body: UserUpdateBody) -> None:
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/user/update",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=NoBody,
|
|
)
|
|
)
|
|
|
|
def delete_user(self, user_id: str) -> None:
|
|
_ = self.proxy.transport.post(
|
|
"/user/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=UserDeleteBody(user_ids=[user_id]),
|
|
response_type=NoBody,
|
|
)
|
|
|
|
def delete_user_strict(self, user_id: str) -> None:
|
|
"""Strict delete for the act phase of a test: a failed delete is a hard
|
|
failure, unlike the warn-only delete_user used at teardown."""
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/user/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=UserDeleteBody(user_ids=[user_id]),
|
|
response_type=UserDeleteResponse,
|
|
)
|
|
)
|
|
|
|
def user_info(self, user_id: str) -> UserInfoResponse:
|
|
return unwrap(
|
|
self.proxy.transport.get(
|
|
"/user/info",
|
|
headers=self.proxy.transport.master,
|
|
params=UserInfoParams(user_id=user_id),
|
|
response_type=UserInfoResponse,
|
|
)
|
|
)
|
|
|
|
def user_count(self, user_id: str) -> int:
|
|
return unwrap(
|
|
self.proxy.transport.get(
|
|
"/user/list",
|
|
headers=self.proxy.transport.master,
|
|
params=UserListParams(user_ids=user_id),
|
|
response_type=UserListResponse,
|
|
)
|
|
).total
|
|
|
|
def user_list_ids(self, user_id: str) -> tuple[str, ...]:
|
|
listing = unwrap(
|
|
self.proxy.transport.get(
|
|
"/user/list",
|
|
headers=self.proxy.transport.master,
|
|
params=UserListParams(user_ids=user_id),
|
|
response_type=UserListResponse,
|
|
)
|
|
)
|
|
return tuple(row.user_id for row in listing.users)
|
|
|
|
def create_org(self, body: OrgNewBody) -> str:
|
|
return unwrap(
|
|
self.proxy.transport.post(
|
|
"/organization/new",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=OrgNewResponse,
|
|
)
|
|
).organization_id
|
|
|
|
def update_org(self, body: OrgUpdateBody) -> None:
|
|
_ = unwrap(
|
|
self.proxy.transport.patch(
|
|
"/organization/update",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=NoBody,
|
|
)
|
|
)
|
|
|
|
def delete_org(self, organization_id: str) -> None:
|
|
_ = self.proxy.transport.delete(
|
|
"/organization/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=OrgDeleteBody(organization_ids=[organization_id]),
|
|
response_type=NoBody,
|
|
)
|
|
|
|
def org_info(self, organization_id: str) -> OrgInfoResponse:
|
|
return unwrap(
|
|
self.proxy.transport.get(
|
|
"/organization/info",
|
|
headers=self.proxy.transport.master,
|
|
params=OrgInfoParams(organization_id=organization_id),
|
|
response_type=OrgInfoResponse,
|
|
)
|
|
)
|
|
|
|
def org_info_status(self, organization_id: str) -> ProbeResult:
|
|
return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id))
|
|
def create_tag(self, body: TagNewBody) -> None:
|
|
_ = unwrap(
|
|
self.proxy.transport.post(
|
|
"/tag/new",
|
|
headers=self.proxy.transport.master,
|
|
json=body,
|
|
response_type=NoBody,
|
|
)
|
|
)
|
|
|
|
def delete_tag(self, name: str) -> None:
|
|
_ = self.proxy.transport.post(
|
|
"/tag/delete",
|
|
headers=self.proxy.transport.master,
|
|
json=TagDeleteBody(name=name),
|
|
response_type=NoBody,
|
|
)
|
|
|
|
def tag_list(self) -> tuple[TagListEntry, ...]:
|
|
return tuple(
|
|
unwrap(
|
|
self.proxy.transport.get(
|
|
"/tag/list",
|
|
headers=self.proxy.transport.master,
|
|
params=NoBody(),
|
|
response_type=TagListResponse,
|
|
)
|
|
).root
|
|
)
|
|
|
|
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
|
|
return self.proxy.transport.send(
|
|
"/chat/completions",
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=ChatBody(model=model, messages=[ChatMessage(role="user", content=content)], max_tokens=16),
|
|
)
|
|
|
|
def key_generate_status(self, key: str, body: KeyGenerateBody) -> StreamingResponse:
|
|
return self.proxy.transport.send("/key/generate", headers=self.proxy.transport.bearer(key), json=body)
|
|
|
|
def team_new_status(self, key: str, body: TeamNewBody) -> StreamingResponse:
|
|
return self.proxy.transport.send("/team/new", headers=self.proxy.transport.bearer(key), json=body)
|
|
|
|
def user_new_status(self, key: str, body: UserNewBody) -> StreamingResponse:
|
|
return self.proxy.transport.send("/user/new", headers=self.proxy.transport.bearer(key), json=body)
|
|
|
|
|
|
def build_client(proxy: ProxyClient) -> ManagementClient:
|
|
return ManagementClient(proxy=proxy, master_key=MASTER_KEY)
|