litellm/tests/integration/authorization/test_team_admin_gate.py

"""Status-code matrix for every management route that admits a team admin today.

Each route is called as a proxy admin, an admin of the target team, a plain member, an admin of another team
and a teamless user. It is called again on a team that belongs to an organization, as an admin of that
organization and as an admin of another organization. The expected codes pin current behaviour so the shared
team-admin gate can prove parity.
"""

from __future__ import annotations

import json
import re
import uuid
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass, replace
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import Final, Literal, assert_never

import httpx
import pytest
from pydantic import JsonValue

from tests.integration._support.client import (
    Gateway,
    Scenario,
    delete_key_if_present,
    eventually,
    gateway_from_environment,
    object_value,
    string_value,
    team_admin_permissions,
)
from tests.integration._support.database import read_rows
from tests.integration._support.wire import Reply, Request, Wire, wire_server

Caller = Literal["proxy_admin", "team_admin", "member", "other_team_admin", "outsider", "org_admin", "other_org_admin"]
CALLERS: Final[tuple[Caller, ...]] = (
    "proxy_admin",
    "team_admin",
    "member",
    "other_team_admin",
    "outsider",
    "org_admin",
    "other_org_admin",
)
ORG_CALLERS: Final[frozenset[Caller]] = frozenset({"org_admin", "other_org_admin"})


@dataclass(frozen=True, slots=True)
class Call:
    method: str
    path: str
    body: Mapping[str, JsonValue] | None = None


@dataclass(frozen=True, slots=True)
class TeamScenario:
    """The shared team as one test case sees it: its ids, a key per caller, and fresh things to act on."""

    scenario: Scenario
    team_id: str
    keys: Mapping[Caller, str]
    request_id: str
    since: datetime
    until: datetime

    @property
    def gateway(self) -> Gateway:
        return self.scenario.gateway

    def user(self) -> str:
        return self.scenario.user(user_role="internal_user")

    def member(self) -> str:
        return self.scenario.member(self.team_id)

    def proxy_admin_member(self) -> str:
        user_id: Final = self.scenario.user(user_role="proxy_admin")
        self.gateway.post("/team/member_add", {"team_id": self.team_id, "member": {"role": "user", "user_id": user_id}})
        return user_id

    def member_key(self) -> str:
        created: Final = self.gateway.post("/key/generate", {"user_id": self.member(), "team_id": self.team_id})
        return string_value(created["key"])

    def service_key(self) -> str:
        created: Final = self.gateway.post(
            "/key/service-account/generate", {"team_id": self.team_id, "key_alias": f"matrix-{uuid.uuid4().hex}"}
        )
        token: Final = string_value(created["key"])
        self.scenario.cleanups.callback(delete_key_if_present, self.gateway, token)
        return token

    def model(self) -> str:
        created: Final = self.gateway.post("/model/new", _team_model_body(self, f"matrix-{uuid.uuid4().hex}"))
        model_id: Final = string_value(object_value(created["model_info"])["id"])
        self.scenario.cleanups.callback(_delete_model_if_present, self.gateway, model_id)
        return model_id

    def callback_name(self) -> str:
        name: Final = f"matrix-{uuid.uuid4().hex}"
        self.scenario.cleanups.callback(self.gateway.request, "DELETE", f"/team/{self.team_id}/callback/{name}")
        return name

    def callback(self) -> str:
        name: Final = self.callback_name()
        self.gateway.post(f"/team/{self.team_id}/callback", _callback_body(name))
        return name

    def invitation(self) -> str:
        created: Final = self.gateway.post("/invitation/new", {"user_id": self.member()}, key=self.keys["team_admin"])
        return string_value(created["id"])


@dataclass(frozen=True, slots=True)
class Route:
    name: str
    call: Callable[[TeamScenario], Call]
    team_admin: int
    others: int
    proxy_admin: int | None = 200
    member: int | None = None
    other_team_admin: int | None = None
    outsider: int | None = None
    org_admin: int | None = None
    other_org_admin: int | None = None
    permission: tuple[str, ...] = ()
    cleanup: Callable[[TeamScenario, dict[str, JsonValue]], None] | None = None

    def expected(self, caller: Caller) -> int | None:
        match caller:
            case "proxy_admin":
                return self.proxy_admin
            case "team_admin":
                return self.team_admin
            case "member":
                return self.others if self.member is None else self.member
            case "other_team_admin":
                return self.others if self.other_team_admin is None else self.other_team_admin
            case "outsider":
                return self.others if self.outsider is None else self.outsider
            case "org_admin":
                return self.others if self.org_admin is None else self.org_admin
            case "other_org_admin":
                return self.others if self.other_org_admin is None else self.other_org_admin
            case _:
                assert_never(caller)


def _day(moment: datetime) -> str:
    return moment.strftime("%Y-%m-%d")


def _stamp(moment: datetime) -> str:
    return moment.strftime("%Y-%m-%d %H:%M:%S")


def _spend_rows(gateway: Gateway, team_id: str, since: datetime, until: datetime) -> list[JsonValue]:
    page: Final = gateway.get(
        "/spend/logs/ui", {"team_id": team_id, "start_date": _stamp(since), "end_date": _stamp(until)}
    )
    rows: Final = page["data"]
    assert isinstance(rows, list)
    return rows


def _team_model_body(s: TeamScenario, name: str, api_base: str | None = None) -> dict[str, JsonValue]:
    return {
        "model_name": name,
        "litellm_params": {
            "model": "openai/gpt-4o-mini",
            "api_key": "integration-provider-key",
            "api_base": api_base or f"{s.gateway.upstream_url}/v1",
        },
        "model_info": {"team_id": s.team_id},
    }


def _callback_body(name: str) -> dict[str, JsonValue]:
    return {
        "callback_name": name,
        "callback_type": "success",
        "callback_vars": {
            "langfuse_public_key": "pk-matrix",
            "langfuse_secret_key": "sk-matrix",
            "langfuse_host": "http://127.0.0.1:9",
        },
    }


def _delete_model_if_present(gateway: Gateway, model_id: str) -> None:
    if read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,)):
        gateway.post("/model/delete", {"id": model_id})


def _delete_project(s: TeamScenario, created: dict[str, JsonValue]) -> None:
    response: Final = s.gateway.request("DELETE", "/project/delete", {"project_ids": [created["project_id"]]})
    assert response.status_code == 200, response.text


def _delete_key(s: TeamScenario, created: dict[str, JsonValue]) -> None:
    delete_key_if_present(s.gateway, string_value(created["key"]))


def _delete_model(s: TeamScenario, created: dict[str, JsonValue]) -> None:
    _delete_model_if_present(s.gateway, string_value(object_value(created["model_info"])["id"]))


# fmt: off
ROUTES: Final[tuple[Route, ...]] = (
    Route("member_add_user",
          lambda s: Call("POST", "/team/member_add", {"team_id": s.team_id, "member": {"role": "user", "user_id": s.user()}}),
          team_admin=200, others=403, org_admin=200),
    Route("member_add_admin",
          lambda s: Call("POST", "/team/member_add", {"team_id": s.team_id, "member": {"role": "admin", "user_id": s.user()}}),
          team_admin=200, others=403, org_admin=200),
    Route("member_update_budget",
          lambda s: Call("POST", "/team/member_update", {"team_id": s.team_id, "user_id": s.member(), "max_budget_in_team": 5}),
          team_admin=200, others=403, org_admin=200),
    Route("member_update_role_admin",
          lambda s: Call("POST", "/team/member_update", {"team_id": s.team_id, "user_id": s.member(), "role": "admin"}),
          team_admin=200, others=403, org_admin=200),
    Route("member_delete",
          lambda s: Call("POST", "/team/member_delete", {"team_id": s.team_id, "user_id": s.member()}),
          team_admin=200, others=403, org_admin=200),
    Route("members_bulk_delete",
          lambda s: Call("POST", f"/management/v1/teams/{s.team_id}/members/bulk_delete", {"members": [{"user_id": s.member()}]}),
          team_admin=200, others=403, org_admin=200),
    Route("members_bulk_update",
          lambda s: Call("POST", f"/management/v1/teams/{s.team_id}/members/bulk_update",
                         {"members": [{"user_id": s.member(), "max_budget_in_team": 10}]}),
          team_admin=200, others=403, org_admin=200),
    Route("member_reset_spend",
          lambda s: Call("POST", f"/team/{s.team_id}/member/{s.member()}/reset_spend", {"reset_to": 0}),
          team_admin=200, others=403, org_admin=200),
    Route("member_reset_budget",
          lambda s: Call("POST", f"/team/{s.team_id}/member/{s.member()}/reset_budget"),
          team_admin=200, others=403, org_admin=200),
    Route("invitation_new",
          lambda s: Call("POST", "/invitation/new", {"user_id": s.member()}),
          team_admin=200, others=400),
    Route("invitation_delete",
          lambda s: Call("POST", "/invitation/delete", {"invitation_id": s.invitation()}),
          team_admin=200, others=400, other_team_admin=403, org_admin=403, other_org_admin=403),
    Route("user_info_v2",
          lambda s: Call("GET", f"/v2/user/info?user_id={s.member()}"),
          team_admin=200, others=404),
    Route("permissions_update",
          lambda s: Call("POST", "/team/permissions_update",
                         {"team_id": s.team_id, "team_member_permissions": ["/key/info", "/key/health"]}),
          team_admin=200, others=403, org_admin=200),
    Route("permissions_list",
          lambda s: Call("GET", f"/team/permissions_list?team_id={s.team_id}"),
          team_admin=200, others=403, org_admin=200),
    Route("key_generate_team",
          lambda s: Call("POST", "/key/generate", {"team_id": s.team_id}),
          team_admin=200, others=400, member=401, cleanup=_delete_key),
    Route("key_generate_for_member",
          lambda s: Call("POST", "/key/generate", {"team_id": s.team_id, "user_id": s.member()}),
          team_admin=403, others=400, member=401, cleanup=_delete_key),
    Route("key_generate_for_member_permitted",
          lambda s: Call("POST", "/key/generate", {"team_id": s.team_id, "user_id": s.member()}),
          team_admin=200, others=400, member=401, permission=("create_member_keys",), cleanup=_delete_key),
    Route("key_generate_for_proxy_admin_member_permitted",
          lambda s: Call("POST", "/key/generate", {"team_id": s.team_id, "user_id": s.proxy_admin_member()}),
          team_admin=403, others=400, member=401, permission=("create_member_keys",), cleanup=_delete_key),
    Route("key_generate_management_key_for_member_permitted",
          lambda s: Call("POST", "/key/generate", {"team_id": s.team_id, "user_id": s.member(), "key_type": "management"}),
          team_admin=403, others=400, member=401, permission=("create_member_keys",), cleanup=_delete_key),
    Route("service_account_generate",
          lambda s: Call("POST", "/key/service-account/generate", {"team_id": s.team_id, "key_alias": f"matrix-{uuid.uuid4().hex}"}),
          team_admin=200, others=400, member=401, cleanup=_delete_key),
    Route("key_update_service_account",
          lambda s: Call("POST", "/key/update", {"key": s.service_key(), "max_budget": 5}),
          team_admin=200, others=401),
    Route("key_update_member_key",
          lambda s: Call("POST", "/key/update", {"key": s.member_key(), "max_budget": 5}),
          team_admin=403, others=403),
    Route("key_update_member_key_permitted",
          lambda s: Call("POST", "/key/update", {"key": s.member_key(), "max_budget": 5}),
          team_admin=200, others=403, permission=("member_key_budgets",)),
    Route("team_key_bulk_update",
          lambda s: Call("POST", "/team/key/bulk_update",
                         {"team_id": s.team_id, "all_keys_in_team": True, "update_fields": {"max_budget": 5}}),
          team_admin=200, others=401),
    Route("key_delete",
          lambda s: Call("POST", "/key/delete", {"keys": [s.member_key()]}),
          team_admin=200, others=403),
    Route("key_regenerate",
          lambda s: Call("POST", "/key/regenerate", {"key": s.member_key()}),
          team_admin=200, others=401),
    Route("key_reset_spend",
          lambda s: Call("POST", f"/key/{s.member_key()}/reset_spend", {"reset_to": 0}),
          team_admin=200, others=403),
    Route("key_block",
          lambda s: Call("POST", "/key/block", {"key": s.member_key()}),
          team_admin=200, others=403, org_admin=200),
    Route("key_unblock",
          lambda s: Call("POST", "/key/unblock", {"key": s.member_key()}),
          team_admin=200, others=403, org_admin=200),
    Route("key_list_team",
          lambda s: Call("GET", f"/key/list?team_id={s.team_id}&include_team_keys=true&return_full_object=true"),
          team_admin=200, others=403, member=200),
    Route("spend_logs_ui",
          lambda s: Call("GET", f"/spend/logs/ui?team_id={s.team_id}&start_date={_stamp(s.since)}&end_date={_stamp(s.until)}"),
          team_admin=200, others=403),
    Route("spend_log_payload",
          lambda s: Call("GET", f"/spend/logs/ui/{s.request_id}"),
          team_admin=200, others=403),
    Route("team_daily_activity",
          lambda s: Call("GET", f"/team/daily/activity?team_ids={s.team_id}&start_date={_day(s.since)}&end_date={_day(s.until)}"),
          team_admin=200, others=404, member=200),
    Route("team_spend_by_user",
          lambda s: Call("GET", f"/team/spend/by_user?team_ids={s.team_id}&start_date={_day(s.since)}&end_date={_day(s.until)}"),
          team_admin=200, others=404, member=200),
    Route("model_new_team",
          lambda s: Call("POST", "/model/new", _team_model_body(s, f"matrix-{uuid.uuid4().hex}")),
          team_admin=200, others=403, cleanup=_delete_model),
    Route("model_update_team",
          lambda s: Call("POST", "/model/update", {"model_info": {"id": s.model(), "team_id": s.team_id}, "litellm_params": {"rpm": 10}}),
          team_admin=200, others=403),
    Route("model_delete_team",
          lambda s: Call("POST", "/model/delete", {"id": s.model()}),
          team_admin=200, others=403),
    Route("auto_router_availability",
          lambda s: Call("POST", "/auto_router/availability", {"team_id": s.team_id}),
          team_admin=200, others=403),
    Route("callback_add",
          lambda s: Call("POST", f"/team/{s.team_id}/callback", _callback_body(s.callback_name())),
          team_admin=200, others=403, org_admin=200),
    Route("callback_get",
          lambda s: Call("GET", f"/team/{s.team_id}/callback"),
          team_admin=200, others=403, org_admin=200),
    Route("callback_delete",
          lambda s: Call("DELETE", f"/team/{s.team_id}/callback/{s.callback()}"),
          team_admin=200, others=403, org_admin=200),
    Route("disable_logging",
          lambda s: Call("POST", f"/team/{s.team_id}/disable_logging"),
          team_admin=401, others=401),
    Route("team_info",
          lambda s: Call("GET", f"/team/info?team_id={s.team_id}"),
          team_admin=200, others=403, member=200, org_admin=200),
    Route("team_update_budget",
          lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 5}),
          team_admin=403, others=403, org_admin=200),
    Route("team_update_budget_permitted",
          lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 4}),
          team_admin=200, others=403, org_admin=200, permission=("max_budget",)),
    Route("team_update_budget_raise",
          lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 6}),
          team_admin=403, others=403, org_admin=200, permission=("max_budget",)),
    Route("team_update_budget_raise_permitted",
          lambda s: Call("POST", "/team/update", {"team_id": s.team_id, "max_budget": 6}),
          team_admin=200, others=403, org_admin=200, permission=("max_budget", "raise_max_budget")),
    Route("project_new",
          lambda s: Call("POST", "/project/new", {"team_id": s.team_id, "project_alias": f"matrix-{uuid.uuid4().hex}"}),
          team_admin=403, others=403, cleanup=_delete_project),
    Route("project_new_permitted",
          lambda s: Call("POST", "/project/new", {"team_id": s.team_id, "project_alias": f"matrix-{uuid.uuid4().hex}"}),
          team_admin=200, others=403, permission=("projects",), cleanup=_delete_project),
    Route("team_delete",
          lambda s: Call("POST", "/team/delete", {"team_ids": [s.team_id]}),
          team_admin=401, others=401, proxy_admin=None),
    Route("team_block",
          lambda s: Call("POST", "/team/block", {"team_id": s.team_id}),
          team_admin=401, others=401, proxy_admin=None),
)
# fmt: on


def _cases() -> Iterator[tuple[Route, Caller]]:
    for route in ROUTES:
        for caller in CALLERS:
            if route.expected(caller) is not None:
                yield route, caller


CASES: Final = tuple(_cases())


def _team_scenario(scenario: Scenario, team_id: str, keys: Mapping[Caller, str]) -> TeamScenario:
    since: Final = datetime.now(timezone.utc) - timedelta(days=1)
    until: Final = since + timedelta(days=2)
    scenario.gateway.chat(scenario.model(), key=keys["team_admin"])
    rows: Final = eventually(
        lambda: _spend_rows(scenario.gateway, team_id, since, until), lambda found: len(found) > 0, seconds=30
    )
    return TeamScenario(
        scenario=scenario,
        team_id=team_id,
        keys=keys,
        request_id=string_value(object_value(rows[0])["request_id"]),
        since=since,
        until=until,
    )


@pytest.fixture(scope="module")
def shared() -> Iterator[TeamScenario]:
    with gateway_from_environment() as gateway, gateway.scenario() as scenario:
        team_id: Final = scenario.team()
        other_team_id: Final = scenario.team()
        team_admin: Final = scenario.member(team_id, role="admin")
        member: Final = scenario.member(team_id)
        other_team_admin: Final = scenario.member(other_team_id, role="admin")
        outsider: Final = scenario.user(user_role="internal_user")
        keys: Final[Mapping[Caller, str]] = MappingProxyType(
            {
                "proxy_admin": gateway.key,
                "team_admin": scenario.key(user_id=team_admin, team_id=team_id),
                "member": scenario.key(user_id=member, team_id=team_id),
                "other_team_admin": scenario.key(user_id=other_team_admin, team_id=other_team_id),
                "outsider": scenario.key(user_id=outsider),
            }
        )
        yield _team_scenario(scenario, team_id, keys)


@pytest.fixture(scope="module")
def org_team() -> Iterator[TeamScenario]:
    with gateway_from_environment() as gateway, gateway.scenario() as scenario:
        organization_id: Final = scenario.organization()
        other_organization_id: Final = scenario.organization()
        team_id: Final = scenario.team(organization_id=organization_id)
        team_admin: Final = scenario.member(team_id, role="admin")
        org_admin: Final = scenario.org_member(organization_id, role="org_admin")
        other_org_admin: Final = scenario.org_member(other_organization_id, role="org_admin")
        keys: Final[Mapping[Caller, str]] = MappingProxyType(
            {
                "team_admin": scenario.key(user_id=team_admin, team_id=team_id),
                "org_admin": scenario.key(user_id=org_admin),
                "other_org_admin": scenario.key(user_id=other_org_admin),
            }
        )
        yield _team_scenario(scenario, team_id, keys)


_BUDGET_BEFORE: Final = 5.0
_BUDGET_ROUTES: Final[Mapping[str, float]] = MappingProxyType(
    {
        "team_update_budget_permitted": 4.0,
        "team_update_budget_raise": 6.0,
        "team_update_budget_raise_permitted": 6.0,
    }
)


@pytest.mark.parametrize(("route", "caller"), CASES, ids=tuple(f"{route.name}[{caller}]" for route, caller in CASES))
def test_status_code(shared: TeamScenario, org_team: TeamScenario, route: Route, caller: Caller) -> None:
    team: Final = org_team if caller in ORG_CALLERS else shared
    with team.gateway.scenario() as scenario:
        s: Final = replace(team, scenario=scenario)
        if route.name in _BUDGET_ROUTES:
            s.gateway.post("/team/update", {"team_id": s.team_id, "max_budget": _BUDGET_BEFORE})
        if route.permission:
            scenario.cleanups.enter_context(team_admin_permissions(s.gateway, route.permission))
        call: Final = route.call(s)
        response: Final = s.gateway.request(call.method, call.path, call.body, key=s.keys[caller])
        assert response.status_code == route.expected(caller), (
            f"{caller} {call.method} {call.path}: {response.status_code} {response.text}"
        )
        if route.name in _BUDGET_ROUTES:
            assert read_rows('SELECT max_budget FROM "LiteLLM_TeamTable" WHERE team_id = %s', (s.team_id,)) == [
                {"max_budget": _BUDGET_ROUTES[route.name] if response.status_code == 200 else _BUDGET_BEFORE}
            ]
        if response.status_code == 200 and route.cleanup is not None:
            route.cleanup(s, object_value(response.json()))


def test_member_key_created_by_a_team_admin_reaches_llm_routes_only(shared: TeamScenario) -> None:
    """The team admin holds the secret of a key that authenticates as the member, so it must not carry the
    member's management access."""
    with shared.gateway.scenario() as scenario:
        s: Final = replace(shared, scenario=scenario)
        scenario.cleanups.enter_context(team_admin_permissions(s.gateway, ("create_member_keys",)))
        member: Final = s.member()
        created: Final = s.gateway.post(
            "/key/generate", {"team_id": s.team_id, "user_id": member}, key=s.keys["team_admin"]
        )
        token: Final = string_value(created["key"])
        scenario.cleanups.callback(delete_key_if_present, s.gateway, token)
        own_key: Final = scenario.key(user_id=member, team_id=s.team_id)

        assert read_rows(
            'SELECT user_id, team_id, allowed_routes FROM "LiteLLM_VerificationToken" WHERE token = %s',
            (string_value(created["token"]),),
        ) == [{"user_id": member, "team_id": s.team_id, "allowed_routes": ["llm_api_routes"]}]
        s.gateway.chat(s.model(), key=token)
        assert s.gateway.request("GET", "/user/info", key=own_key).status_code == 200
        assert s.gateway.request("GET", "/user/info", key=token).status_code == 403


ADMIN_ONLY_CALLERS: Final[tuple[tuple[str, Caller], ...]] = (
    ("shared", "team_admin"),
    ("shared", "member"),
    ("shared", "outsider"),
    ("org_team", "org_admin"),
)
ADMIN_ONLY_IDS: Final = tuple(caller for _, caller in ADMIN_ONLY_CALLERS)


def _admin_only_team(request: pytest.FixtureRequest, fixture: str, scenario: Scenario) -> TeamScenario:
    team: Final = request.getfixturevalue(fixture)
    assert isinstance(team, TeamScenario)
    return replace(team, scenario=scenario)


def _team_model(s: TeamScenario, api_base: str | None = None) -> tuple[str, str]:
    name: Final = f"matrix-{uuid.uuid4().hex}"
    created: Final = s.gateway.post("/model/new", _team_model_body(s, name, api_base))
    model_id: Final = string_value(object_value(created["model_info"])["id"])
    s.scenario.cleanups.callback(_delete_model_if_present, s.gateway, model_id)
    return model_id, name


def _blocked(model_id: str) -> list[dict[str, object]]:
    return read_rows('SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,))


def _serve(s: TeamScenario, name: str, key: str) -> httpx.Response:
    return s.gateway.request(
        "POST",
        "/v1/chat/completions",
        {"model": name, "messages": [{"role": "user", "content": "serving state"}]},
        key=key,
    )


def _is_model_discovery(request: Request) -> bool:
    return (request.method, request.target) == ("GET", "/v1/models")


def _model_serving_reply(request: Request) -> Reply:
    if _is_model_discovery(request):
        return Reply(body=b'{"object":"list","data":[]}')
    assert (request.method, request.target) == ("POST", "/v1/chat/completions"), request
    return Reply(
        body=json.dumps(
            {
                "id": "chatcmpl-model-blocked-control",
                "object": "chat.completion",
                "created": 0,
                "model": "gpt-4o-mini",
                "choices": [
                    {
                        "index": 0,
                        "message": {"role": "assistant", "content": "model serving control"},
                        "finish_reason": "stop",
                    }
                ],
                "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
            }
        ).encode()
    )


def _provider_calls(wire: Wire) -> tuple[Request, ...]:
    return tuple(entry for entry in wire.drain() if not _is_model_discovery(entry))


def _assert_serving_state(s: TeamScenario, name: str, wire: Wire, *, blocked: bool) -> None:
    served: Final = _serve(s, name, s.keys["team_admin"])
    if blocked:
        assert (served.status_code, served.json()) == _MODEL_BLOCKED_REFUSED, served.text
        assert _provider_calls(wire) == (), f"Blocked model reached the provider: {served.text}"
        return

    assert served.status_code == 200, served.text
    assert served.json()["choices"][0]["message"]["content"] == "model serving control", served.text
    received: Final = _provider_calls(wire)
    assert [
        (entry.method, entry.target, entry.headers["authorization"], json.loads(entry.body)) for entry in received
    ] == [
        (
            "POST",
            "/v1/chat/completions",
            "Bearer integration-provider-key",
            {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "serving state"}]},
        )
    ], served.text


def _refusal(response: httpx.Response) -> JsonValue:
    """The error body with the caller's masked user id elided, since it differs per caller."""
    body: Final = response.json()
    error: Final = object_value(body).get("error")
    if isinstance(error, dict) and isinstance(error.get("message"), str):
        return {
            **body,
            "error": {**error, "message": re.sub(r"Your user_id=\S+", "Your user_id=<caller>", error["message"])},
        }
    return body


def _admin_only_route(route: str) -> tuple[int, JsonValue]:
    return 401, {
        "error": {
            "message": "Authentication Error, Only proxy admin can be used to generate, delete, update info for new "
            f"keys/users/teams. Route={route}. Your role=internal_user. Your user_id=<caller>",
            "type": "auth_error",
            "param": "None",
            "code": "401",
        }
    }


_BLOCK_REFUSALS: Final[Mapping[tuple[str, str], tuple[int, JsonValue]]] = MappingProxyType(
    {
        ("POST", "/model/block"): _admin_only_route("/model/block"),
        ("POST", "/model/unblock"): _admin_only_route("/model/unblock"),
    }
)
_PATCH_BLOCKED_REFUSED: Final[tuple[int, JsonValue]] = (
    403,
    {
        "error": {
            "message": "Only proxy admins can change a model's blocked flag.",
            "type": "auth_error",
            "param": "blocked",
            "code": "403",
        }
    },
)
_MODEL_BLOCKED_REFUSED: Final[tuple[int, JsonValue]] = (
    403,
    {
        "error": {
            "message": "litellm.PermissionDeniedError: Model is blocked",
            "type": "permission_error",
            "param": None,
            "code": "403",
        }
    },
)


@pytest.mark.parametrize(("fixture", "caller"), ADMIN_ONLY_CALLERS, ids=ADMIN_ONLY_IDS)
def test_only_proxy_admin_flips_a_models_blocked_flag(
    request: pytest.FixtureRequest, fixture: str, caller: Caller
) -> None:
    with (
        gateway_from_environment() as gateway,
        gateway.scenario() as scenario,
        wire_server(_model_serving_reply) as wire,
    ):
        s: Final = _admin_only_team(request, fixture, scenario)
        model_id, name = _team_model(s, f"{wire.url}/v1")
        internal_key: Final = (
            s.scenario.key(user_id=s.scenario.user(user_role="internal_user")) if caller == "team_admin" else None
        )
        for blocked in (True, False):
            s.gateway.post("/model/block" if blocked else "/model/unblock", {"model_id": model_id})
            attempts = (
                Call("POST", "/model/unblock" if blocked else "/model/block", {"model_id": model_id}),
                Call("PATCH", f"/model/{model_id}/update", {"blocked": not blocked}),
            )
            for call in attempts:
                response = s.gateway.request(call.method, call.path, call.body, key=s.keys[caller])
                if call.method == "PATCH" and caller != "team_admin":
                    assert response.status_code == 403, f"{caller} {call.method} {call.path}: {response.text}"
                else:
                    expected = (
                        _PATCH_BLOCKED_REFUSED if call.method == "PATCH" else _BLOCK_REFUSALS[call.method, call.path]
                    )
                    assert (response.status_code, _refusal(response)) == expected, (
                        f"{caller} {call.method} {call.path}: {response.text}"
                    )
                assert _blocked(model_id) == [{"blocked": blocked}], (
                    f"{caller} {call.method} {call.path} flipped blocked"
                )
            if internal_key is not None:
                path = "/model/unblock" if blocked else "/model/block"
                internal_response = s.gateway.request("POST", path, {"model_id": model_id}, key=internal_key)
                assert (internal_response.status_code, _refusal(internal_response)) == _BLOCK_REFUSALS["POST", path], (
                    internal_response.text
                )
                assert _blocked(model_id) == [{"blocked": blocked}], internal_response.text
            _assert_serving_state(s, name, wire, blocked=blocked)


def test_a_key_allowed_the_block_routes_still_cannot_flip_a_models_blocked_flag(shared: TeamScenario) -> None:
    with wire_server(_model_serving_reply) as wire:
        user: Final = shared.scenario.user(user_role="internal_user")
        key: Final = shared.scenario.key(user_id=user, allowed_routes=["/model/block", "/model/unblock"])
        model_id, name = _team_model(shared, f"{wire.url}/v1")
        for blocked in (True, False):
            shared.gateway.post("/model/block" if blocked else "/model/unblock", {"model_id": model_id})
            call = Call("POST", "/model/unblock" if blocked else "/model/block", {"model_id": model_id})
            response = shared.gateway.request(call.method, call.path, call.body, key=key)
            assert (response.status_code, response.json()) == _PATCH_BLOCKED_REFUSED, response.text
            assert _blocked(model_id) == [{"blocked": blocked}], response.text
            _assert_serving_state(shared, name, wire, blocked=blocked)


@contextmanager
def _model_creation_disabled_for_internal_users(gateway: Gateway) -> Iterator[None]:
    setting: Final = "disable_model_add_for_internal_users"
    original: Final = object_value(gateway.get("/get/ui_settings")["values"]).get(setting, False)
    response: Final = gateway.request("PATCH", "/update/ui_settings", {setting: True})
    assert response.status_code == 200, response.text
    try:
        yield
    finally:
        restored: Final = gateway.request("PATCH", "/update/ui_settings", {setting: original})
        assert restored.status_code == 200, restored.text


def _model_rows(name: str) -> list[dict[str, object]]:
    return read_rows(
        'SELECT model_id FROM "LiteLLM_ProxyModelTable" '
        "WHERE model_name = %s OR model_info->>'team_public_model_name' = %s",
        (name, name),
    )


def _listed_names(gateway: Gateway, key: str) -> tuple[frozenset[str], frozenset[str]]:
    info: Final = gateway.get("/model/info")["data"]
    models: Final = gateway.request("GET", "/v1/models", key=key)
    assert models.status_code == 200, models.text
    listed: Final = models.json()["data"]
    assert isinstance(info, list) and isinstance(listed, list), models.text
    info_names: Final = frozenset(
        name
        for entry in map(object_value, info)
        for name in (entry["model_name"], object_value(entry["model_info"]).get("team_public_model_name"))
        if isinstance(name, str)
    )
    return info_names, frozenset(string_value(object_value(entry)["id"]) for entry in listed)


def test_internal_user_model_creation_prohibition_refuses_team_admin_but_not_proxy_admin(
    shared: TeamScenario, peer: Gateway
) -> None:
    with shared.gateway.scenario() as scenario:
        s: Final = replace(shared, scenario=scenario)
        team_admin: Final = s.keys["team_admin"]
        refused_name: Final = f"matrix-{uuid.uuid4().hex}"
        admin_name: Final = f"matrix-{uuid.uuid4().hex}"
        with _model_creation_disabled_for_internal_users(peer):
            refused: Final = s.gateway.request("POST", "/model/new", _team_model_body(s, refused_name), key=team_admin)
            assert (refused.status_code, refused.json()) == (
                403,
                {
                    "error": {
                        "message": "Model creation is disabled for internal users by disable_model_add_for_internal_users.",
                        "type": "auth_error",
                        "param": "None",
                        "code": "403",
                    }
                },
            ), refused.text
            assert _model_rows(refused_name) == []
            info_names, listed = _listed_names(s.gateway, team_admin)
            assert refused_name not in info_names and refused_name not in listed, (info_names, listed)
            created: Final = s.gateway.post("/model/new", _team_model_body(s, admin_name))
            s.scenario.cleanups.callback(
                _delete_model_if_present, s.gateway, string_value(object_value(created["model_info"])["id"])
            )
            assert _model_rows(admin_name) == [{"model_id": object_value(created["model_info"])["id"]}]
            assert admin_name in _listed_names(s.gateway, team_admin)[1]
        allowed_name: Final = f"matrix-{uuid.uuid4().hex}"
        allowed: Final = s.gateway.request("POST", "/model/new", _team_model_body(s, allowed_name), key=team_admin)
        assert allowed.status_code == 200, allowed.text
        s.scenario.cleanups.callback(
            _delete_model_if_present, s.gateway, string_value(object_value(allowed.json()["model_info"])["id"])
        )
        assert len(_model_rows(allowed_name)) == 1


def _credential_row(name: str) -> list[dict[str, object]]:
    return read_rows(
        'SELECT credential_name, credential_values, credential_info, updated_at::text AS updated_at FROM "LiteLLM_CredentialsTable" '
        "WHERE credential_name = %s",
        (name,),
    )


@pytest.mark.parametrize("contains_slash", (False, True), ids=("plain-name", "slash-name"))
@pytest.mark.parametrize(("fixture", "caller"), ADMIN_ONLY_CALLERS, ids=ADMIN_ONLY_IDS)
def test_only_proxy_admin_reads_or_changes_credentials(
    request: pytest.FixtureRequest, fixture: str, caller: Caller, contains_slash: bool
) -> None:
    secret: Final = f"synthetic-shared-credential-{uuid.uuid4().hex}"
    with gateway_from_environment() as gateway, gateway.scenario() as scenario:
        s: Final = _admin_only_team(request, fixture, scenario)
        name: Final = f"credential/{uuid.uuid4().hex}" if contains_slash else f"credential-{uuid.uuid4().hex}"
        attempted: Final = f"credential/{uuid.uuid4().hex}" if contains_slash else f"credential-{uuid.uuid4().hex}"
        s.gateway.post(
            "/credentials",
            {"credential_name": name, "credential_values": {"api_key": secret}, "credential_info": {"team": "shared"}},
        )
        s.scenario.cleanups.callback(s.gateway.request, "DELETE", f"/credentials/{name}")
        s.scenario.cleanups.callback(s.gateway.request, "DELETE", f"/credentials/{attempted}")
        model_name: Final = f"credential-model-{uuid.uuid4().hex}"
        created_model: Final = s.gateway.post(
            "/model/new",
            {
                "model_name": model_name,
                "litellm_params": {
                    "model": "openai/gpt-4o-mini",
                    "api_base": f"{s.gateway.upstream_url}/v1",
                    "litellm_credential_name": name,
                },
            },
        )
        model_id: Final = string_value(object_value(created_model["model_info"])["id"])
        s.scenario.cleanups.callback(_delete_model_if_present, s.gateway, model_id)
        credential_for_model: Final = s.gateway.request("GET", f"/credentials/by_model/{model_id}")
        assert credential_for_model.status_code == 200, credential_for_model.text
        assert credential_for_model.json() == {
            "credential_name": f"{model_name}-credential-{model_id}",
            "display_name": None,
            "credential_values": {"api_base": f"{s.gateway.upstream_url}/v1"},
            "credential_info": {},
        }, credential_for_model.text
        stored: Final = _credential_row(name)
        assert len(stored) == 1, stored
        attempts: Final = (
            Call(
                "POST",
                "/credentials",
                {"credential_name": attempted, "credential_values": {"api_key": "k"}, "credential_info": {}},
            ),
            Call(
                "PATCH",
                f"/credentials/{name}",
                {"credential_name": name, "credential_values": {"api_key": "overwritten"}, "credential_info": {}},
            ),
            Call("DELETE", f"/credentials/{name}"),
            Call("GET", "/credentials"),
            Call("GET", f"/credentials/by_name/{name}"),
            Call("GET", f"/credentials/by_model/{model_id}"),
        )
        for call in attempts:
            response = s.gateway.request(call.method, call.path, call.body, key=s.keys[caller])
            assert (response.status_code, _refusal(response)) == _admin_only_route(call.path), (
                f"{caller} {call.method} {call.path}: {response.text}"
            )
            assert secret not in response.text, f"{caller} {call.method} {call.path} exposed the credential"
            assert _credential_row(name) == stored, f"{caller} {call.method} {call.path} changed the credential row"
            assert _credential_row(attempted) == [], f"{caller} {call.method} {call.path} wrote a credential"


@pytest.mark.parametrize(
    ("fixture", "caller"),
    (("shared", "member"), ("shared", "outsider"), ("org_team", "org_admin")),
    ids=("member", "outsider", "org_admin"),
)
def test_non_team_admin_patching_the_blocked_flag_gets_the_blocked_flag_refusal(
    request: pytest.FixtureRequest, fixture: str, caller: Caller
) -> None:
    pytest.skip(
        'BUG: PATCH /model/{id}/update {"blocked": ...} by a non-team-admin returns the auto-router refusal '
        "instead of the blocked-flag refusal (https://github.com/BerriAI/litellm/issues/44792)"
    )
    with gateway_from_environment() as gateway, gateway.scenario() as scenario:
        s: Final = _admin_only_team(request, fixture, scenario)
        model_id, _ = _team_model(s)
        response: Final = s.gateway.request("PATCH", f"/model/{model_id}/update", {"blocked": True}, key=s.keys[caller])
        assert (response.status_code, _refusal(response)) == _PATCH_BLOCKED_REFUSED, response.text
        assert _blocked(model_id) == [{"blocked": False}], response.text