mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
4c3879a88e
Slice 4 of the management-endpoints behavior-pinning effort. New ``actors.py`` defines the actor enum + seeds an immutable world (2 orgs, 2 teams, 8 users, 8 verification tokens) under the ``behavior-pin-`` prefix so the rows are identifiable in psql and ``_wipe_world`` is targeted. Each actor key is created with its cleartext form generated locally and its hashed form (via ``litellm.proxy.utils.hash_token``) stored in ``LiteLLM_VerificationToken`` — so the real ``user_api_key_auth`` accepts the cleartext bearer token. Roles, ``team_id``, ``organization_id``, and the service-account metadata flag are all set on the seeded rows so the auth layer resolves the same scopes a real proxy would. The session-scoped ``world`` fixture re-seeds at session start (idempotent via wipe-then-create), and the smoke test confirms each of the 8 actor keys can call ``/key/info`` on itself and receive its own row back. Plan: https://www.notion.so/36643b8acdab8128a581ced0f6a4744d
273 lines
9 KiB
Python
273 lines
9 KiB
Python
"""Immutable read-world seed for behavior-pinning tests.
|
|
|
|
Seeds 8 actor profiles across 2 orgs / 2 teams so authz matrix tests can drive
|
|
``(actor, target, expected)`` tuples against the real proxy. Each actor has
|
|
exactly one virtual key (cleartext + hashed) so tests pass the cleartext as a
|
|
``Bearer`` token and the real auth stack accepts it.
|
|
|
|
All seeded rows are namespaced under the ``behavior-pin-`` prefix on their
|
|
primary key, so the seed is identifiable in psql and the wipe step is targeted.
|
|
"""
|
|
|
|
import enum
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict
|
|
|
|
from prisma import Json
|
|
|
|
from litellm.proxy._types import LitellmUserRoles
|
|
from litellm.proxy.utils import PrismaClient, hash_token
|
|
|
|
|
|
class Actor(str, enum.Enum):
|
|
PROXY_ADMIN = "proxy_admin"
|
|
ORG_ADMIN = "org_admin"
|
|
TEAM_ADMIN = "team_admin"
|
|
INTERNAL_USER = "internal_user"
|
|
OWNER = "owner"
|
|
UNRELATED_SAME_ORG = "unrelated_same_org"
|
|
CROSS_ORG_USER = "cross_org_user"
|
|
SERVICE_ACCOUNT = "service_account"
|
|
|
|
|
|
# Stable IDs so re-seed is idempotent and the world is identifiable in psql.
|
|
PREFIX = "behavior-pin-"
|
|
ORG_A = PREFIX + "org-a"
|
|
ORG_B = PREFIX + "org-b"
|
|
TEAM_ALPHA = PREFIX + "team-alpha"
|
|
TEAM_BETA = PREFIX + "team-beta"
|
|
BUDGET_ID = PREFIX + "budget"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SeededKey:
|
|
user_id: str
|
|
cleartext: str
|
|
hashed: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class World:
|
|
org_a_id: str
|
|
org_b_id: str
|
|
team_alpha_id: str
|
|
team_beta_id: str
|
|
keys: Dict[Actor, SeededKey]
|
|
|
|
|
|
def _new_clear_key() -> str:
|
|
return "sk-" + uuid.uuid4().hex
|
|
|
|
|
|
def _actor_profile() -> Dict[Actor, Dict[str, Any]]:
|
|
"""Per-actor (role, scoping) profile used for both the user row and its key.
|
|
|
|
Scoping rules:
|
|
- PROXY_ADMIN: global, no team/org scope on its key.
|
|
- ORG_ADMIN: org_a, no team scope.
|
|
- TEAM_ADMIN / INTERNAL_USER / OWNER / UNRELATED_SAME_ORG / SERVICE_ACCOUNT:
|
|
team_alpha within org_a.
|
|
- CROSS_ORG_USER: team_beta within org_b.
|
|
|
|
The auth layer reads ``user_role`` off the user row pointed at by the key's
|
|
``user_id``, so setting it once on the user is enough — keys do not need a
|
|
separate role field.
|
|
"""
|
|
return {
|
|
Actor.PROXY_ADMIN: {
|
|
"user_role": LitellmUserRoles.PROXY_ADMIN.value,
|
|
"team_id": None,
|
|
"organization_id": None,
|
|
},
|
|
Actor.ORG_ADMIN: {
|
|
"user_role": LitellmUserRoles.ORG_ADMIN.value,
|
|
"team_id": None,
|
|
"organization_id": ORG_A,
|
|
},
|
|
Actor.TEAM_ADMIN: {
|
|
"user_role": LitellmUserRoles.INTERNAL_USER.value,
|
|
"team_id": TEAM_ALPHA,
|
|
"organization_id": ORG_A,
|
|
},
|
|
Actor.INTERNAL_USER: {
|
|
"user_role": LitellmUserRoles.INTERNAL_USER.value,
|
|
"team_id": TEAM_ALPHA,
|
|
"organization_id": ORG_A,
|
|
},
|
|
Actor.OWNER: {
|
|
"user_role": LitellmUserRoles.INTERNAL_USER.value,
|
|
"team_id": TEAM_ALPHA,
|
|
"organization_id": ORG_A,
|
|
},
|
|
Actor.UNRELATED_SAME_ORG: {
|
|
"user_role": LitellmUserRoles.INTERNAL_USER.value,
|
|
"team_id": TEAM_ALPHA,
|
|
"organization_id": ORG_A,
|
|
},
|
|
Actor.CROSS_ORG_USER: {
|
|
"user_role": LitellmUserRoles.INTERNAL_USER.value,
|
|
"team_id": TEAM_BETA,
|
|
"organization_id": ORG_B,
|
|
},
|
|
Actor.SERVICE_ACCOUNT: {
|
|
"user_role": LitellmUserRoles.INTERNAL_USER.value,
|
|
"team_id": TEAM_ALPHA,
|
|
"organization_id": ORG_A,
|
|
},
|
|
}
|
|
|
|
|
|
async def _wipe_world(prisma: PrismaClient) -> None:
|
|
"""Delete prior seed rows so re-seed is idempotent across sessions."""
|
|
await prisma.db.litellm_verificationtoken.delete_many(
|
|
where={"user_id": {"startswith": PREFIX}}
|
|
)
|
|
await prisma.db.litellm_organizationmembership.delete_many(
|
|
where={"user_id": {"startswith": PREFIX}}
|
|
)
|
|
await prisma.db.litellm_teammembership.delete_many(
|
|
where={"user_id": {"startswith": PREFIX}}
|
|
)
|
|
await prisma.db.litellm_usertable.delete_many(
|
|
where={"user_id": {"startswith": PREFIX}}
|
|
)
|
|
await prisma.db.litellm_teamtable.delete_many(
|
|
where={"team_id": {"startswith": PREFIX}}
|
|
)
|
|
await prisma.db.litellm_organizationtable.delete_many(
|
|
where={"organization_id": {"startswith": PREFIX}}
|
|
)
|
|
await prisma.db.litellm_budgettable.delete_many(where={"budget_id": BUDGET_ID})
|
|
|
|
|
|
async def seed_world(prisma: PrismaClient) -> World:
|
|
await _wipe_world(prisma)
|
|
|
|
# Budget — orgs require a non-null budget_id.
|
|
await prisma.db.litellm_budgettable.create(
|
|
data={
|
|
"budget_id": BUDGET_ID,
|
|
"created_by": "behavior-pin-seeder",
|
|
"updated_by": "behavior-pin-seeder",
|
|
}
|
|
)
|
|
|
|
# Orgs.
|
|
for org_id, alias in [(ORG_A, "alpha"), (ORG_B, "beta")]:
|
|
await prisma.db.litellm_organizationtable.create(
|
|
data={
|
|
"organization_id": org_id,
|
|
"organization_alias": alias,
|
|
"budget_id": BUDGET_ID,
|
|
"created_by": "behavior-pin-seeder",
|
|
"updated_by": "behavior-pin-seeder",
|
|
}
|
|
)
|
|
|
|
profiles = _actor_profile()
|
|
user_ids: Dict[Actor, str] = {actor: PREFIX + actor.value for actor in Actor}
|
|
|
|
# Users.
|
|
for actor, profile in profiles.items():
|
|
teams_list = [profile["team_id"]] if profile["team_id"] else []
|
|
await prisma.db.litellm_usertable.create(
|
|
data={
|
|
"user_id": user_ids[actor],
|
|
"user_role": profile["user_role"],
|
|
"team_id": profile["team_id"],
|
|
"organization_id": profile["organization_id"],
|
|
"teams": teams_list,
|
|
}
|
|
)
|
|
|
|
# Teams.
|
|
await prisma.db.litellm_teamtable.create(
|
|
data={
|
|
"team_id": TEAM_ALPHA,
|
|
"team_alias": "alpha-1",
|
|
"organization_id": ORG_A,
|
|
"admins": [user_ids[Actor.TEAM_ADMIN]],
|
|
"members": [
|
|
user_ids[Actor.TEAM_ADMIN],
|
|
user_ids[Actor.INTERNAL_USER],
|
|
user_ids[Actor.OWNER],
|
|
user_ids[Actor.UNRELATED_SAME_ORG],
|
|
user_ids[Actor.SERVICE_ACCOUNT],
|
|
],
|
|
}
|
|
)
|
|
await prisma.db.litellm_teamtable.create(
|
|
data={
|
|
"team_id": TEAM_BETA,
|
|
"team_alias": "beta-1",
|
|
"organization_id": ORG_B,
|
|
"admins": [],
|
|
"members": [user_ids[Actor.CROSS_ORG_USER]],
|
|
}
|
|
)
|
|
|
|
# Org memberships (UI lookups + ORG_ADMIN scoping use these).
|
|
for actor, org_id, role in [
|
|
(Actor.ORG_ADMIN, ORG_A, "org_admin"),
|
|
(Actor.TEAM_ADMIN, ORG_A, "internal_user"),
|
|
(Actor.INTERNAL_USER, ORG_A, "internal_user"),
|
|
(Actor.OWNER, ORG_A, "internal_user"),
|
|
(Actor.UNRELATED_SAME_ORG, ORG_A, "internal_user"),
|
|
(Actor.SERVICE_ACCOUNT, ORG_A, "internal_user"),
|
|
(Actor.CROSS_ORG_USER, ORG_B, "internal_user"),
|
|
]:
|
|
await prisma.db.litellm_organizationmembership.create(
|
|
data={
|
|
"user_id": user_ids[actor],
|
|
"organization_id": org_id,
|
|
"user_role": role,
|
|
}
|
|
)
|
|
|
|
# Team memberships.
|
|
for actor, team_id in [
|
|
(Actor.TEAM_ADMIN, TEAM_ALPHA),
|
|
(Actor.INTERNAL_USER, TEAM_ALPHA),
|
|
(Actor.OWNER, TEAM_ALPHA),
|
|
(Actor.UNRELATED_SAME_ORG, TEAM_ALPHA),
|
|
(Actor.SERVICE_ACCOUNT, TEAM_ALPHA),
|
|
(Actor.CROSS_ORG_USER, TEAM_BETA),
|
|
]:
|
|
await prisma.db.litellm_teammembership.create(
|
|
data={"user_id": user_ids[actor], "team_id": team_id}
|
|
)
|
|
|
|
# Verification tokens — one cleartext per actor, hashed via the real
|
|
# credential boundary so user_api_key_auth accepts it.
|
|
keys: Dict[Actor, SeededKey] = {}
|
|
for actor, profile in profiles.items():
|
|
cleartext = _new_clear_key()
|
|
hashed = hash_token(cleartext)
|
|
token_data: Dict[str, Any] = {
|
|
"token": hashed,
|
|
"key_name": PREFIX + actor.value + "-key",
|
|
"user_id": user_ids[actor],
|
|
# LiteLLM_VerificationTokenView is non-Optional on models — Postgres lets the
|
|
# column be NULL, but the pydantic view used by user_api_key_auth rejects None.
|
|
"models": [],
|
|
}
|
|
if profile["team_id"]:
|
|
token_data["team_id"] = profile["team_id"]
|
|
if profile["organization_id"]:
|
|
token_data["organization_id"] = profile["organization_id"]
|
|
if actor == Actor.SERVICE_ACCOUNT:
|
|
# LiteLLM convention: service-account keys carry an explicit metadata flag.
|
|
token_data["metadata"] = Json({"service_account_id": user_ids[actor]})
|
|
await prisma.db.litellm_verificationtoken.create(data=token_data)
|
|
keys[actor] = SeededKey(
|
|
user_id=user_ids[actor], cleartext=cleartext, hashed=hashed
|
|
)
|
|
|
|
return World(
|
|
org_a_id=ORG_A,
|
|
org_b_id=ORG_B,
|
|
team_alpha_id=TEAM_ALPHA,
|
|
team_beta_id=TEAM_BETA,
|
|
keys=keys,
|
|
)
|