From 109ca70f668dcef80ec71b281697a5d559854a2c Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 19:51:29 -0700 Subject: [PATCH] feat(auto-router): allow opted-in team members to manage their routers --- litellm/proxy/_types.py | 2 + litellm/proxy/auth/auth_checks.py | 16 +- litellm/proxy/auth/auto_router_checks.py | 136 +++++++ litellm/proxy/auth/user_api_key_auth.py | 2 +- .../common_utils/encrypt_decrypt_utils.py | 2 +- .../auto_router_endpoints.py | 117 ++++-- .../model_management_endpoints.py | 273 ++++++++++++-- .../management_endpoints/team_endpoints.py | 37 +- .../auto_router_permissions.py | 345 ++++++++++++++++++ litellm/repositories/prisma_protocols.py | 5 + litellm/router.py | 47 ++- litellm/types/router.py | 1 + .../proxy/auth/test_auth_checks.py | 35 +- .../test_auto_router_endpoints.py | 127 ++++++- .../test_model_management_endpoints.py | 235 +++++++++++- .../test_team_endpoints.py | 18 + .../test_auto_router_permissions.py | 208 +++++++++++ tests/test_litellm/test_router.py | 251 ++++++++++++- .../test_router_model_cost_isolation.py | 11 +- .../AutoRouters/AutoRoutersPanel.tsx | 1 + .../components/AutoRouters/autoRouterRows.ts | 13 +- .../(dashboard)/models-and-endpoints/page.tsx | 15 +- .../panels/AutoRoutersTabPanel.test.tsx | 23 +- .../panels/AutoRoutersTabPanel.tsx | 7 +- .../add_model/add_auto_router_tab.test.tsx | 61 +++- .../add_model/add_auto_router_tab.tsx | 41 ++- .../build_auto_router_test_targets.ts | 56 +++ .../handle_add_auto_router_submit.tsx | 12 +- .../common_components/team_dropdown.test.tsx | 47 ++- .../common_components/team_dropdown.tsx | 33 +- ...dit_auto_router_modal.integration.test.tsx | 56 ++- .../edit_auto_router_modal.tsx | 55 ++- .../components/key_team_helpers/key_list.tsx | 1 + .../llm_calls/fetch_models.test.tsx | 21 +- .../src/components/llm_calls/fetch_models.tsx | 13 + .../src/components/model_info_view.tsx | 73 +--- .../team/permission_definitions.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 11 +- .../src/utils/modelPermissions.test.ts | 31 +- .../src/utils/modelPermissions.ts | 41 ++- 40 files changed, 2254 insertions(+), 226 deletions(-) create mode 100644 litellm/proxy/auth/auto_router_checks.py create mode 100644 litellm/proxy/management_helpers/auto_router_permissions.py create mode 100644 tests/test_litellm/proxy/management_helpers/test_auto_router_permissions.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ffc41a9d7ae..c88f7a84e35 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -284,6 +284,7 @@ class KeyManagementRoutes(str, enum.Enum): # team's `team_member_permissions`, non-admin members of that team may set # `access_group_ids` on keys they create/update. Default-deny. KEY_ACCESS_GROUP_ASSIGNMENT = "/key/access_group_assignment" + AUTO_ROUTER_MANAGE = "/auto_router/manage" # info and health routes KEY_INFO = "/key/info" @@ -650,6 +651,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.KEY_RESET_SPEND.value, KeyManagementRoutes.KEY_ALIASES.value, KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value, + KeyManagementRoutes.AUTO_ROUTER_MANAGE.value, ] management_routes = ( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6d61ad4d3e8..a90d087a27b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -39,6 +39,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.models.project import LiteLLM_ProjectTable from litellm.proxy._types import ( RBAC_ROLES, CallInfo, @@ -109,7 +110,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.organization_repository import OrganizationRepository -from litellm.repositories.prisma_protocols import RowT_co +from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( AccessGroupRepository, @@ -847,6 +848,7 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset( "/health", "/health/services", "/health/test_connection", + "/auto_router/test_routing", } ) @@ -3172,7 +3174,7 @@ async def _delete_cache_access_object( @log_db_metrics async def get_access_object( access_group_id: str, - prisma_client: PrismaClient | None, + prisma_client: DatabaseClient | None, user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_AccessGroupTable: @@ -3918,7 +3920,7 @@ async def get_org_object( async def _get_resources_from_access_groups( access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], - prisma_client: PrismaClient | None = None, + prisma_client: DatabaseClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> list[str]: @@ -3976,7 +3978,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( access_group_ids: Sequence[str], - prisma_client: PrismaClient | None = None, + prisma_client: DatabaseClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, ) -> list[str]: @@ -4475,6 +4477,7 @@ async def can_key_call_model( llm_model_list: Sequence[object] | None, valid_token: UserAPIKeyAuth, llm_router: litellm.Router | None, + prisma_client: DatabaseClient | None = None, ) -> Literal[True]: """ Checks if token can call a given model @@ -4504,6 +4507,7 @@ async def can_key_call_model( if key_access_group_ids: models_from_groups: Final = await _get_models_from_access_groups( access_group_ids=key_access_group_ids, + prisma_client=prisma_client, ) if models_from_groups: return _can_object_call_model( @@ -4632,6 +4636,7 @@ async def can_team_access_model( team_object: LiteLLM_TeamTable | None, llm_router: Router | None, team_model_aliases: dict[str, str] | None = None, + prisma_client: DatabaseClient | None = None, ) -> Literal[True]: """ Returns True if the team can access a specific model. @@ -4654,6 +4659,7 @@ async def can_team_access_model( if team_access_group_ids: models_from_groups: Final = await _get_models_from_access_groups( access_group_ids=team_access_group_ids, + prisma_client=prisma_client, ) if models_from_groups: return _can_object_call_model( @@ -4749,7 +4755,7 @@ async def _key_access_group_grants_model( def can_project_access_model( model: str | list[str], - project_object: LiteLLM_ProjectTableCachedObj, + project_object: LiteLLM_ProjectTable, llm_router: Router | None, ) -> Literal[True]: """ diff --git a/litellm/proxy/auth/auto_router_checks.py b/litellm/proxy/auth/auto_router_checks.py new file mode 100644 index 00000000000..b83e1f3fffe --- /dev/null +++ b/litellm/proxy/auth/auto_router_checks.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + +if TYPE_CHECKING: + from litellm.router import Router + +_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _mapping(value: object) -> Mapping[str, object] | None: + try: + return _MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +async def authorize_member_auto_router_inference( + *, + deployment: Mapping[str, object] | None, + request_kwargs: Mapping[str, object], + llm_router: Router, +) -> None: + if deployment is None: + return + model_info: Final = _mapping(deployment.get("model_info")) + if model_info is None or model_info.get("member_auto_router") is not True: + return + + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import ( + OrganizationNotFoundError, + TeamNotFoundError, + get_org_object, + get_project_object, + get_team_membership, + get_team_object, + ) + from litellm.proxy.management_helpers.auto_router_permissions import ( + MemberAutoRouterDependencyObjects, + authorize_member_auto_router_dependencies, + validate_member_auto_router_config, + ) + + metadata: Final = _mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs))) + actor: Final = metadata.get("user_api_key_auth") if metadata is not None else None + team_id: Final = model_info.get("team_id") + if not isinstance(actor, UserAPIKeyAuth) or not isinstance(team_id, str) or not team_id: + raise HTTPException(status_code=403, detail="Member auto-routers require authenticated team access") + if actor.team_id != team_id and actor.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException(status_code=403, detail="This auto-router belongs to a different team") + + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + raise HTTPException(status_code=503, detail="Cannot verify auto-router model access without a database") + try: + team: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=actor.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except TeamNotFoundError as error: + raise HTTPException(status_code=403, detail="The auto-router team no longer exists") from error + if ( + actor.user_role != LitellmUserRoles.PROXY_ADMIN + and actor.user_id is not None + and (not actor.user_id or not any(member.user_id == actor.user_id for member in team.members_with_roles)) + ): + raise HTTPException(status_code=403, detail="You are no longer a member of this auto-router's team") + if team.blocked: + raise HTTPException(status_code=403, detail="This auto router's team is blocked.") + params: Final = _mapping(deployment.get("litellm_params")) + if params is None: + raise HTTPException(status_code=403, detail="The member auto-router configuration is invalid") + raw_config: Final = _mapping(params.get("complexity_router_config")) + if raw_config is None: + raise HTTPException(status_code=403, detail="The member auto-router configuration is invalid") + default_model: Final = params.get("complexity_router_default_model") + config: Final = validate_member_auto_router_config(raw_config) + membership: Final = ( + await get_team_membership( + user_id=actor.user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=actor.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if actor.user_id + else None + ) + try: + organization: Final = ( + await get_org_object( + org_id=team.organization_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=actor.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if team.organization_id + else None + ) + except OrganizationNotFoundError as error: + raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.") from error + project: Final = ( + await get_project_object( + project_id=actor.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if actor.project_id + else None + ) + await authorize_member_auto_router_dependencies( + config=config, + default_model=default_model if isinstance(default_model, str) else None, + user_api_key_dict=actor, + team=team, + prisma_client=None, + llm_router=llm_router, + dependency_objects=MemberAutoRouterDependencyObjects( + membership=membership, organization=organization, project=project + ), + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 25570ab220a..0cf5259c1f8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2489,7 +2489,7 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, request: Request, - request_data: dict, + request_data: dict[str, object], route: str, ) -> None: """Run ``common_checks`` once at the ``user_api_key_auth`` wrapper diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index fd9b3beee46..288dedebbc6 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -124,7 +124,7 @@ def decrypt_value_helper( key: str, # this is just for debug purposes, showing the k,v pair that's invalid. not a signing key. exception_type: Literal["debug", "error"] = "error", return_original_value: bool = False, -): +) -> str | None: signing_key: Final = _get_salt_key() try: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 50716e5d474..200ed6c3bf3 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -40,6 +40,14 @@ from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, refresh_proxy_server_request_body_snapshot, ) +from litellm.proxy.management_endpoints.common_utils import ( + _is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership +) +from litellm.proxy.management_helpers.auto_router_permissions import ( + authorize_member_auto_router_dependencies, + authorize_member_auto_router_team, + validate_member_auto_router_config, +) from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.team_repository import TeamRepository @@ -72,13 +80,13 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( ) if TYPE_CHECKING: - from fastapi import APIRouter, Depends, HTTPException, Query, status + from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm.proxy.utils import PrismaClient from litellm.router import Router else: try: - from fastapi import APIRouter, Depends, HTTPException, Query, status + from fastapi import APIRouter, Depends, HTTPException, Query, Request, status except ImportError: # fastapi is only required for proxy, not for SDK usage pass @@ -201,21 +209,14 @@ async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) - return await prisma_client.db.query_raw(query, *args) -async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> None: - """Allow exactly the callers who could create this router. - - Both dry runs are gated like the write they rehearse rather than as reads: a proxy - admin, or a team admin naming their own team, matching /model/new. Routing a test - prompt can also spend money (an `llm` classifier config calls its classifier, a - semantic config embeds the prompt), so a read-level gate would be too loose anyway. - """ +async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> LiteLLM_TeamTable | None: from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, ) from litellm.proxy.proxy_server import premium_user, prisma_client if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: - return + return None if team_id is None: raise HTTPException( @@ -244,12 +245,47 @@ async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: }, ) - ModelManagementAuthChecks.can_user_make_team_model_call( - team_id=team_id, + team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump()) + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team): + ModelManagementAuthChecks.can_user_make_team_model_call( + team_id=team_id, + user_api_key_dict=user_api_key_dict, + team_obj=team, + premium_user=premium_user, + ) + return None + authorize_member_auto_router_team( user_api_key_dict=user_api_key_dict, - team_obj=LiteLLM_TeamTable.model_validate(team_row.model_dump()), + team=team, premium_user=premium_user, ) + return team + + +async def _authorize_member_dry_run_config( + *, + config: Mapping[str, object], + default_model: str | None, + user_api_key_dict: UserAPIKeyAuth, + team: LiteLLM_TeamTable, +) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import llm_router, prisma_client + + if prisma_client is None or llm_router is None: + raise HTTPException(status_code=503, detail="Cannot verify auto-router model access") + validated: Final = validate_member_auto_router_config(config) + scoped_actor: Final = user_api_key_dict.model_copy( + update=MappingProxyType({"team_id": team.team_id, "team_models": team.models, "org_id": team.organization_id}) + ) + await authorize_member_auto_router_dependencies( + config=validated, + default_model=default_model, + user_api_key_dict=scoped_actor, + team=team, + prisma_client=prisma_client, + llm_router=llm_router, + ) + return scoped_actor def _models_this_test_can_call(config: RequestComplexityRouterConfig) -> tuple[str, ...]: @@ -326,16 +362,23 @@ async def validate_complexity_router_config( Runs the same check every write path runs (the router's own pydantic model), so a form can show the backend's exact verdict while the operator is still editing rather than after a - rejected save. Gated exactly like the save it rehearses: a proxy admin, or a team admin - naming their own team. Nothing is created, routed, or billed. + rejected save. Uses the same team opt-in and model-access checks as configuration + writes for members. Nothing is created, routed, or billed. """ - await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, ) error: Final = validate_complexity_router_config_write(data.complexity_router_config) + if error is None and member_team is not None: + await _authorize_member_dry_run_config( + config=data.complexity_router_config, + default_model=None, + user_api_key_dict=user_api_key_dict, + team=member_team, + ) return ComplexityRouterConfigValidationResponse(valid=error is None, error=error) @@ -349,6 +392,7 @@ async def validate_complexity_router_config( async def preview_auto_router_routing( data: AutoRouterRoutingTestRequest, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + http_request: Request, ) -> AutoRouterRoutingTestResponse: """ Route a single request through a complexity-router config and report where it landed. @@ -392,7 +436,34 @@ async def preview_auto_router_routing( ) from litellm.proxy.utils import get_available_models_for_user - await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id) + actor: Final = ( + await _authorize_member_dry_run_config( + config=data.complexity_router_config.model_dump(exclude_none=True), + default_model=data.default_model, + user_api_key_dict=user_api_key_dict, + team=member_team, + ) + if member_team is not None + else user_api_key_dict + ) + request_data: Final[dict[str, object]] = { # mutable-ok: auth and routing enrich this request in place + **data.wire_body(), + "metadata": {}, # mutable-ok: centralized auth and identity stamping share this metadata bucket + "proxy_server_request": {"body": None}, # mutable-ok: the snapshot owner fills this body in place + } + + if member_team is not None and _models_this_test_can_call(data.complexity_router_config): + from litellm.proxy.auth.user_api_key_auth import ( + _run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse the serving admission policy + ) + + await _run_centralized_common_checks( + user_api_key_auth_obj=actor, + request=http_request, + request_data=request_data, + route="/auto_router/test_routing", + ) if llm_router is None: raise HTTPException( @@ -404,7 +475,7 @@ async def preview_auto_router_routing( await _authorize_models_this_test_can_call( config=data.complexity_router_config, - user_api_key_dict=user_api_key_dict, + user_api_key_dict=actor, llm_router=llm_router, ) @@ -417,12 +488,8 @@ async def preview_auto_router_routing( ) request_kwargs: Final = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data={ # mutable-ok: the request-metadata helper takes and returns request kwargs as a dict - **data.wire_body(), - "metadata": {}, # mutable-ok: the request-metadata helper writes the auth fields into this dict - "proxy_server_request": {"body": None}, # mutable-ok: the snapshot owner fills body in place - }, - user_api_key_dict=user_api_key_dict, + data=request_data, + user_api_key_dict=actor, _metadata_variable_name="metadata", ) refresh_proxy_server_request_body_snapshot(request_kwargs) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 2234e825090..bcddb1f7ef0 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -15,13 +15,16 @@ import datetime import json from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass +from fnmatch import fnmatchcase from json import JSONDecodeError from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast, runtime_checkable from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator +import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME @@ -51,6 +54,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY +from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -65,6 +69,7 @@ from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( _refresh_cached_team, + append_team_models, team_model_add, team_model_delete, ) @@ -76,6 +81,13 @@ from litellm.proxy.management_helpers.access_group_model_sync import ( sync_access_groups_for_renamed_model, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.management_helpers.auto_router_permissions import ( + MemberAutoRouterWrite, + StoredAutoRouterIdentity, + authorize_member_auto_router_dependencies, + authorize_member_auto_router_team, + authorize_member_auto_router_write, +) from litellm.proxy.spend_tracking.ptu_feature_flag import ( PTU_COST_ATTRIBUTION_ENV_VAR, is_ptu_cost_attribution_enabled, @@ -122,12 +134,14 @@ from litellm.types.router import ( GenericLiteLLMParams, ModelInfo, updateDeployment, + updateLiteLLMParams, ) from litellm.types.utils import without_server_derived_pricing from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import models as prisma_models + from prisma import types as prisma_types router: Final = APIRouter() @@ -180,6 +194,24 @@ class _ProxyModelTable(Protocol): class _TxModelTables(Protocol): litellm_proxymodeltable: _ProxyModelTable + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + +@runtime_checkable +class _TransactionFactory(Protocol): + def __call__(self, *, timeout: datetime.timedelta = ...) -> AbstractAsyncContextManager[_TxModelTables]: ... + + +class _ModelTransactionClient(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True) + + tx: _TransactionFactory + + +@dataclass(frozen=True, slots=True) +class _TransactionClient: + db: _TxModelTables + _RowT = TypeVar("_RowT") @@ -213,7 +245,7 @@ def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable: def _repo_team_table(prisma_client: PrismaClient) -> _TeamLookupTable: - return TeamRepository(prisma_client).table + return TeamRepository(WriterPinnedClient(prisma_client.db)).table def _db_team_table(prisma_client: PrismaClient) -> _TeamTable: @@ -353,6 +385,25 @@ def _effective_complexity_router_params( ) +def _member_auto_router_marker_for_update( + *, + incoming_params: updateLiteLLMParams | None, + existing: Deployment, + member_write: MemberAutoRouterWrite | None, +) -> bool | None: + if member_write is not None: + return True + if not existing.model_info.member_auto_router: + return None + if incoming_params is None: + return True + if any(getattr(incoming_params, field, None) is not None for field in STRATEGY_ROUTER_PARAM_FIELDS): + return False + if incoming_params.model is not None and incoming_params.model != _effective_model(None, existing.litellm_params): + return False + return True + + def _decrypted_model(stored_model: object) -> str | None: if not isinstance(stored_model, str): return None @@ -385,7 +436,11 @@ def _raise_on_tuning_quota_violation( @asynccontextmanager async def _auto_router_capability_slot( - prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None + prisma_client: PrismaClient, + *, + effective_params: Mapping[str, object], + model_id: str | None, + member_write: MemberAutoRouterWrite | None = None, ) -> AsyncGenerator[_ProxyModelTable, None]: """Hand out the model table to write through while the row's claim on a licensed capability is settled. @@ -394,9 +449,8 @@ async def _auto_router_capability_slot( (a statement's snapshot predates anything it locks), so pods cannot both pass the count: the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged against the license limit and the write is refused with a 403 before it happens. The row - being edited keeps its own slot through ``model_id``. Every other write, and every write on - an unlimited license, goes through the repository table with no lock. Only the row write - itself may run inside: anything that needs a second connection (the team model bookkeeping) + being edited keeps its own slot through ``model_id``. Member writes also recheck their + authorization under this lock. Team model bookkeeping needs a second connection and must wait until the transaction has committed and the lock is released. The transaction writes bypass the repository's publish-on-write, so the config change is published once after commit, the way delete_team_models does. @@ -408,6 +462,7 @@ async def _auto_router_capability_slot( _license_check, # pyright: ignore[reportPrivateUsage] # existing capability slot reads the proxy license singleton heuristic_v1_tuning_baselines, llm_router, + premium_user, ) limit: Final = _license_check.auto_router_capability_limit() @@ -415,13 +470,96 @@ async def _auto_router_capability_slot( baselines: Final = heuristic_v1_tuning_baselines tuning_candidate: Final = _tuning_candidate(effective_params, model_id=model_id) judges_tuning: Final = baselines is not None and is_mutable_tuned_candidate(tuning_candidate, baselines) - if limit is None or (capability is None and not judges_tuning): + if member_write is None and (limit is None or (capability is None and not judges_tuning)): yield _proxy_model_table(prisma_client) return - async with prisma_client.db.tx() as tx_ctx: + transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) + transaction: Final = ( + transaction_client.tx(timeout=datetime.timedelta(seconds=30)) + if member_write is not None + else transaction_client.tx() + ) + async with transaction as tx_ctx: tables: Final[_TxModelTables] = tx_ctx await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + if member_write is not None: + if member_write.model_id is not None: + await tx_ctx.query_raw( + 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE', + member_write.model_id, + ) + pinned_client: Final = _TransactionClient(tx_ctx) + team_where: Final[prisma_types.LiteLLM_TeamTableWhereUniqueInput] = {"team_id": member_write.team_id} + team_include: Final[prisma_types.LiteLLM_TeamTableInclude] = {"litellm_model_table": True} + team_row: Final = await TeamRepository(pinned_client).table.find_unique( + where=team_where, include=team_include + ) + if team_row is None or llm_router is None: + raise HTTPException(status_code=403, detail="The auto router's team or model catalog is unavailable.") + team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump()) + authorize_member_auto_router_team( + user_api_key_dict=member_write.actor, team=team, premium_user=premium_user + ) + if member_write.model_id is not None: + model_where: Final[prisma_types.LiteLLM_ProxyModelTableWhereInput] = {"model_id": member_write.model_id} + current_row: Final = await tables.litellm_proxymodeltable.find_unique(where=model_where) + current_identity: Final = ( + StoredAutoRouterIdentity.model_validate(current_row.model_dump()) + if current_row is not None + else None + ) + current_model: Final = ( + Deployment.model_validate(current_row.model_dump()) if current_row is not None else None + ) + if ( + current_identity is None + or current_identity.created_by != member_write.actor.user_id + or current_model is None + or current_model.model_info.team_id != member_write.team_id + ): + raise HTTPException(status_code=403, detail="Team members can update only their own auto routers.") + if current_identity.updated_at != member_write.updated_at: + raise HTTPException(status_code=409, detail="This auto router changed. Reload it before updating.") + else: + all_models: Final[prisma_types.LiteLLM_ProxyModelTableWhereInput] = {} + rows_for_names: Final = await tables.litellm_proxymodeltable.find_many(where=all_models) + stored_names: Final = tuple( + ( + row.model_name, + model_info_as_mapping(row.model_info), + ) + for row in rows_for_names + ) + config_names: Final = tuple( + (str(row.get("model_name", "")), model_info_as_mapping(row.get("model_info"))) + for row in config_rows + ) + team_aliases: Final = team_model_aliases(team) + aliases: Final = ( + *(llm_router.model_group_alias or ()), + *(litellm.model_alias_map or ()), + *(team_aliases or ()), + ) + if member_write.public_name in aliases or any( + fnmatchcase( + member_write.public_name, + str(info.get("team_public_model_name") or name) + if info is not None and info.get("team_id") == member_write.team_id + else name, + ) + for name, info in (*stored_names, *config_names) + if info is None or info.get("team_id") in (None, member_write.team_id) + ): + raise HTTPException(status_code=409, detail="This auto-router name is already used by a model.") + await authorize_member_auto_router_dependencies( + config=member_write.config, + default_model=member_write.default_model, + user_api_key_dict=member_write.actor, + team=team, + prisma_client=pinned_client, + llm_router=llm_router, + ) if capability is not None: rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" @@ -434,7 +572,7 @@ async def _auto_router_capability_slot( status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" ) if judges_tuning and baselines is not None: - model_rows: Final = await ModelRepository(WriterPinnedClient(tx_ctx)).find_all_except(model_id or "") + model_rows: Final = await ModelRepository(_TransactionClient(tx_ctx)).find_all_except(model_id or "") _raise_on_tuning_quota_violation( candidate=tuning_candidate, others=tuple( @@ -883,11 +1021,39 @@ async def patch_model( param=None, ) - await ModelManagementAuthChecks.can_user_make_model_call( + write_authorization: Final = await ModelManagementAuthChecks.can_user_make_model_call( model_params=db_model, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + member_operation="update", + incoming_model_params=patch_data, + ) + member_write: Final = write_authorization if isinstance(write_authorization, MemberAutoRouterWrite) else None + member_marker: Final = _member_auto_router_marker_for_update( + incoming_params=patch_data.litellm_params, existing=db_model, member_write=member_write + ) + marker_info: Final = ( + ModelInfo(id=db_model.model_info.id) + if member_write is not None + else patch_data.model_info or ModelInfo(id=db_model.model_info.id) + ) + effective_info: Final = ( + marker_info.model_copy(update=MappingProxyType({"member_auto_router": member_marker})) + if member_marker is not None + else patch_data.model_info + ) + effective_patch: Final = ( + patch_data.model_copy( + update=MappingProxyType( + { + "model_name": None if member_write is not None else patch_data.model_name, + "model_info": effective_info, + } + ) + ) + if member_marker is not None + else patch_data ) # Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins @@ -933,13 +1099,14 @@ async def patch_model( prisma_client, effective_params=effective_params, model_id=model_id, + member_write=member_write, ) as table: return await table.update(where={"model_id": model_id}, data=update_data) # Handle team model updates with proper alias management updated_model: Final = await _update_team_model_in_db( db_model=db_model, - patch_data=patch_data, + patch_data=effective_patch, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, write_row=write_row, @@ -1218,7 +1385,7 @@ async def _add_team_model_to_db( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, -) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable | None": """ If 'team_id' is provided, @@ -1226,6 +1393,8 @@ async def _add_team_model_to_db( - store the model in the db with the unique 'model_name' - add the public model name to the team's allowed models list """ + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + _team_id: Final = model_params.model_info.team_id if _team_id is None: return None @@ -1253,13 +1422,14 @@ async def _add_team_model_to_db( ) if original_model_name: - await team_model_add( + await append_team_models( data=TeamModelAddRequest( team_id=_team_id, models=[original_model_name], ), - http_request=Request(scope={"type": "http"}), - user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, ) return model_response @@ -1787,9 +1957,16 @@ class ModelManagementAuthChecks: prisma_client: PrismaClient, premium_user: bool, allow_missing_team: bool = False, - ) -> Literal[True]: + member_operation: Literal["create", "update"] | None = None, + incoming_model_params: updateDeployment | None = None, + ) -> Literal[True] | MemberAutoRouterWrite: + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ): + raise HTTPException(status_code=403, detail="View-only users cannot manage models.") ## Check team model auth - if model_params.model_info is not None and model_params.model_info.team_id is not None: + if model_params.model_info.team_id is not None: team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1810,6 +1987,27 @@ class ModelManagementAuthChecks: ) team_obj: Final = LiteLLM_TeamTable.model_validate(team_obj_row.model_dump()) + if ( + member_operation is not None + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + and not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + ): + from litellm.proxy.proxy_server import llm_router + + if llm_router is None or (member_operation == "update" and incoming_model_params is None): + raise HTTPException( + status_code=400, detail="An auto-router configuration and model catalog are required." + ) + return await authorize_member_auto_router_write( + incoming=incoming_model_params if incoming_model_params is not None else model_params, + existing=model_params if member_operation == "update" else None, + user_api_key_dict=user_api_key_dict, + team=team_obj, + premium_user=premium_user, + prisma_client=prisma_client, + llm_router=llm_router, + ) + return ModelManagementAuthChecks.can_user_make_team_model_call( team_id=model_params.model_info.team_id, user_api_key_dict=user_api_key_dict, @@ -2067,12 +2265,14 @@ async def add_new_model( ) ## Auth check - await ModelManagementAuthChecks.can_user_make_model_call( + write_authorization: Final = await ModelManagementAuthChecks.can_user_make_model_call( model_params=model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + member_operation="create", ) + member_write: Final = write_authorization if isinstance(write_authorization, MemberAutoRouterWrite) else None ModelManagementAuthChecks.can_user_attach_credential( litellm_params=model_params.litellm_params, @@ -2094,9 +2294,14 @@ async def add_new_model( enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)), ) - model_params.model_info = ModelInfo( # rebind-ok: downstream team-model handling mutates this same object + clean_model_info: Final = ModelInfo( **without_server_derived_pricing(model_params.model_info.model_dump(exclude_none=True)) ) + model_params.model_info = ( # rebind-ok: downstream team-model handling mutates this same object + clean_model_info.model_copy(update=MappingProxyType({"member_auto_router": True})) + if member_write is not None + else clean_model_info + ) model_response: prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None = None # update DB @@ -2129,6 +2334,7 @@ async def add_new_model( None, ), model_id=priced_model_params.model_info.id, + member_write=member_write, ), ) reload_outcome = await proxy_config.add_deployment( @@ -2259,12 +2465,15 @@ async def update_model( raise Exception("model not found") deployment: Final = Deployment(**_existing_litellm_params.model_dump()) - await ModelManagementAuthChecks.can_user_make_model_call( + write_authorization: Final = await ModelManagementAuthChecks.can_user_make_model_call( model_params=deployment, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + member_operation="update", + incoming_model_params=model_params, ) + member_write: Final = write_authorization if isinstance(write_authorization, MemberAutoRouterWrite) else None ModelManagementAuthChecks.can_user_attach_credential( litellm_params=model_params.litellm_params, @@ -2285,6 +2494,9 @@ async def update_model( effective_params: Final = _effective_complexity_router_params( model_params.litellm_params, deployment.litellm_params ) + member_marker: Final = _member_auto_router_marker_for_update( + incoming_params=model_params.litellm_params, existing=deployment, member_write=member_write + ) # update DB if store_model_in_db is True: @@ -2317,15 +2529,30 @@ async def update_model( and deployment.model_info.team_id is None else None ) - _data: Final[dict[str, str]] = { + base_update: Final[PrismaCompatibleUpdateDBModel] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, - **({} if renamed_to is None else {"model_name": renamed_to}), } + renamed_update: Final[PrismaCompatibleUpdateDBModel] = ( + {**base_update, "model_name": renamed_to} # mutable-ok: Prisma serializes only concrete update dicts + if renamed_to is not None + else base_update + ) + _data: Final[PrismaCompatibleUpdateDBModel] = ( + { # mutable-ok: Prisma serializes only concrete update dicts + **renamed_update, + "model_info": deployment.model_info.model_copy( + update=MappingProxyType({"member_auto_router": member_marker}) + ).model_dump_json(exclude_none=True), + } + if member_marker is not None + else renamed_update + ) async with _auto_router_capability_slot( prisma_client, effective_params=effective_params, model_id=_model_id, + member_write=member_write, ) as table: model_response: Final = await table.update( where={"model_id": _model_id}, @@ -2421,7 +2648,6 @@ async def update_public_model_groups( """ try: # Update the public model groups - import litellm from litellm.proxy.proxy_server import proxy_config, store_model_in_db # Check if user has admin permissions @@ -2496,7 +2722,6 @@ async def update_useful_links( """ try: # Update the public model groups - import litellm from litellm.proxy.proxy_server import proxy_config # Check if user has admin permissions diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 9350d2cd691..c52eeeff2d8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3309,7 +3309,8 @@ async def team_member_delete( }' ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3447,6 +3448,25 @@ async def team_member_delete( } ) + await delete_cache_team_object( + team_id=data.team_id, + team_alias=existing_team_row.team_alias, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await delete_cache_key_objects( + hashed_tokens=tuple(key.token for key in keys_to_delete), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await evict_and_broadcast(cache_keys=tuple(sorted(user_ids_to_delete)), user_api_key_cache=user_api_key_cache) + for user_id in sorted(user_ids_to_delete): + await invalidate_team_member_spend_state( + user_id=user_id, + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) + _emit_team_members_metric(existing_team_row) return existing_team_row @@ -5668,6 +5688,21 @@ async def team_model_add( detail={"error": "Only proxy admin or team admin can modify team models"}, ) + return await append_team_models( + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def append_team_models( + *, + data: TeamModelAddRequest, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> "prisma_models.LiteLLM_TeamTable": # Atomic array append with dedup at the database level so concurrent # BYOK model creates don't overwrite each other's team.models entries. # When the team currently has models=[] (unrestricted access), the diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py new file mode 100644 index 00000000000..381c966f2f0 --- /dev/null +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -0,0 +1,345 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.project import LiteLLM_ProjectTable +from litellm.proxy._types import ( + UI_TEAM_ID, + CommonProxyErrors, + KeyManagementRoutes, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LitellmUserRoles, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _check_team_member_model_access, # pyright: ignore[reportPrivateUsage] # shared membership authorization owner + can_key_call_model, + can_org_access_model, + can_project_access_model, + can_team_access_model, +) +from litellm.proxy.auth.team_grants import team_model_aliases +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import TeamMembershipRepository +from litellm.router import Router +from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model, strategy_router_dependencies +from litellm.types.management_endpoints.auto_router_endpoints import RequestComplexityRouterConfig +from litellm.types.router import Deployment, updateDeployment + +if TYPE_CHECKING: + from prisma import types as prisma_types + + +class _MemberRouterThinking(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + type: Literal["enabled", "disabled", "adaptive"] + budget_tokens: int | None = Field(default=None, gt=0, le=1_000_000) + + +class _MemberRouterGenerationParams(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + reasoning_effort: str | None = None + thinking: _MemberRouterThinking | None = None + verbosity: Literal["low", "medium", "high"] | None = None + max_tokens: int | None = Field(default=None, gt=0, le=1_000_000) + max_completion_tokens: int | None = Field(default=None, gt=0, le=1_000_000) + max_output_tokens: int | None = Field(default=None, gt=0, le=1_000_000) + temperature: float | None = Field(default=None, ge=0, le=2, allow_inf_nan=False) + top_p: float | None = Field(default=None, ge=0, le=1, allow_inf_nan=False) + frequency_penalty: float | None = Field(default=None, ge=-2, le=2, allow_inf_nan=False) + presence_penalty: float | None = Field(default=None, ge=-2, le=2, allow_inf_nan=False) + seed: int | None = None + stop: str | tuple[str, ...] | None = None + + +class _MemberComplexityRouterConfig(RequestComplexityRouterConfig): + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + +class _RouterConfigSource(BaseModel): + model: str | None = None + complexity_router_config: Mapping[str, object] | None = None + + +class _MembershipKey(TypedDict): + user_id: ReadOnly[str] + team_id: ReadOnly[str] + + +class _MembershipWhere(TypedDict): + user_id_team_id: ReadOnly[_MembershipKey] + + +@dataclass(frozen=True, slots=True) +class MemberAutoRouterDependencyObjects: + membership: LiteLLM_TeamMembership | None + organization: LiteLLM_OrganizationTable | None + project: LiteLLM_ProjectTable | None + + +def authorize_member_auto_router_team( + *, user_api_key_dict: UserAPIKeyAuth, team: LiteLLM_TeamTable, premium_user: bool +) -> None: + if not premium_user: + raise HTTPException(status_code=403, detail=CommonProxyErrors.not_premium_user.value) + if ( + user_api_key_dict.user_role + not in (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.TEAM, LitellmUserRoles.ORG_ADMIN) + or not user_api_key_dict.user_id + or not any(member.user_id == user_api_key_dict.user_id for member in team.members_with_roles) + or user_api_key_dict.team_id not in (None, UI_TEAM_ID, team.team_id) + or team.blocked + or KeyManagementRoutes.AUTO_ROUTER_MANAGE.value not in (team.team_member_permissions or ()) + ): + raise HTTPException(status_code=403, detail="This team does not allow you to manage your own auto routers.") + + +def validate_member_auto_router_config(config: Mapping[str, object]) -> RequestComplexityRouterConfig: + try: + validated: Final = _MemberComplexityRouterConfig.model_validate(config) + for entries in validated.tier_model_configs.values(): + for entry in entries: + _MemberRouterGenerationParams.model_validate(entry.litellm_params) + return validated + except ValidationError as exc: + location: Final = ".".join(str(part) for part in exc.errors()[0]["loc"]) + raise HTTPException(status_code=400, detail=f"Invalid member auto-router configuration at {location}.") from exc + + +async def authorize_member_auto_router_dependencies( + *, + config: RequestComplexityRouterConfig, + default_model: str | None, + user_api_key_dict: UserAPIKeyAuth, + team: LiteLLM_TeamTable, + prisma_client: DatabaseClient | None, + llm_router: Router, + dependency_objects: MemberAutoRouterDependencyObjects | None = None, +) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + if team.blocked: + raise HTTPException(status_code=403, detail="This auto router's team is blocked.") + aliases: Final = team_model_aliases(team) + alias_dict: Final = ( + dict(aliases) if aliases is not None else None # mutable-ok: auth model and helpers require dict + ) + scoped_actor: Final = user_api_key_dict.model_copy( + update=MappingProxyType({"team_id": team.team_id, "team_models": team.models, "team_model_aliases": alias_dict}) + ) + objects: Final = ( + dependency_objects + if dependency_objects is not None + else await _load_member_auto_router_dependency_objects( + user_api_key_dict=scoped_actor, team=team, prisma_client=prisma_client + ) + ) + if team.organization_id and objects.organization is None: + raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.") + if scoped_actor.project_id and ( + objects.project is None or objects.project.team_id != team.team_id or objects.project.blocked + ): + raise HTTPException(status_code=403, detail="The auto router's project is unavailable.") + dependencies: Final = strategy_router_dependencies( + MappingProxyType( + { + "model": "auto_router/complexity_router", + "complexity_router_config": config.model_dump(exclude_none=True), + "complexity_router_default_model": default_model, + } + ) + ) + for model, deployments in ( + (dependency.model_name, llm_router.get_model_list(model_name=dependency.model_name, team_id=team.team_id)) + for dependency in dependencies + ): + if not deployments or any( + classify_strategy_router_model(_RouterConfigSource.model_validate(deployment["litellm_params"]).model or "") + is not None + for deployment in deployments + ): + raise HTTPException(status_code=400, detail=f"Auto-router target {model!r} must be a configured model.") + await can_team_access_model( + model=model, + team_object=team, + llm_router=llm_router, + team_model_aliases=alias_dict, + prisma_client=prisma_client, + ) + await can_key_call_model( + model=model, + llm_model_list=None, + valid_token=scoped_actor, + llm_router=llm_router, + prisma_client=prisma_client, + ) + await _check_team_member_model_access( + model=model, + team_object=team, + valid_token=scoped_actor, + llm_router=llm_router, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=objects.membership, + team_membership_loaded=True, + ) + if objects.organization is not None: + can_org_access_model(model=model, org_object=objects.organization, llm_router=llm_router) + if objects.project is not None: + can_project_access_model(model=model, project_object=objects.project, llm_router=llm_router) + + +async def _load_member_auto_router_dependency_objects( + *, user_api_key_dict: UserAPIKeyAuth, team: LiteLLM_TeamTable, prisma_client: DatabaseClient | None +) -> MemberAutoRouterDependencyObjects: + if prisma_client is None: + raise HTTPException(status_code=503, detail="Cannot verify auto-router model access without a database") + membership_where: Final[_MembershipWhere] = { + "user_id_team_id": {"user_id": user_api_key_dict.user_id or "", "team_id": team.team_id} + } + membership_include: Final[prisma_types.LiteLLM_TeamMembershipInclude] = {"litellm_budget_table": True} + membership_row: Final = ( + await TeamMembershipRepository(prisma_client).table.find_unique( + where=membership_where, include=membership_include + ) + if user_api_key_dict.user_id + else None + ) + membership: Final = ( + LiteLLM_TeamMembership.model_validate(membership_row.model_dump()) if membership_row is not None else None + ) + organization: Final = ( + await OrganizationRepository(prisma_client).find_by_id(team.organization_id) if team.organization_id else None + ) + if team.organization_id and organization is None: + raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.") + project: Final = ( + await ProjectRepository(prisma_client).find_by_id(user_api_key_dict.project_id) + if user_api_key_dict.project_id + else None + ) + return MemberAutoRouterDependencyObjects(membership=membership, organization=organization, project=project) + + +class StoredAutoRouterIdentity(BaseModel): + created_by: str | None = None + updated_at: datetime | None = None + + +@dataclass(frozen=True, slots=True) +class MemberAutoRouterWrite: + actor: UserAPIKeyAuth + team_id: str + model_id: str | None + public_name: str + updated_at: datetime | None + config: RequestComplexityRouterConfig + default_model: str | None + + +async def authorize_member_auto_router_write( + *, + incoming: Deployment | updateDeployment, + existing: Deployment | None, + user_api_key_dict: UserAPIKeyAuth, + team: LiteLLM_TeamTable, + premium_user: bool, + prisma_client: DatabaseClient, + llm_router: Router, +) -> MemberAutoRouterWrite: + authorize_member_auto_router_team(user_api_key_dict=user_api_key_dict, team=team, premium_user=premium_user) + stored: Final = StoredAutoRouterIdentity.model_validate(existing.model_dump()) if existing is not None else None + if stored is not None and stored.created_by != user_api_key_dict.user_id: + raise HTTPException(status_code=403, detail="Team members can update only their own auto routers.") + params: Final = incoming.litellm_params + if params is None or incoming.model_fields_set - frozenset({"model_name", "litellm_params", "model_info"}): + raise HTTPException(status_code=403, detail="Team members may change only auto-router configuration.") + if params.model_fields_set - frozenset({"model", "complexity_router_config", "complexity_router_default_model"}): + raise HTTPException(status_code=403, detail="Team members may change only auto-router configuration.") + info: Final = incoming.model_info + if info is not None and ( + info.model_fields_set - frozenset({"id", "team_id"}) + or info.team_id not in (None, team.team_id) + or (existing is not None and "id" in info.model_fields_set and info.id != existing.model_info.id) + ): + raise HTTPException( + status_code=403, detail="Team members cannot change model ownership or administrative settings." + ) + existing_model: Final = ( + decrypt_value_helper(existing.litellm_params.model, key="model", return_original_value=True) + if existing is not None + else None + ) + effective_model: Final = params.model or existing_model + if ( + not isinstance(effective_model, str) + or classify_strategy_router_model(effective_model) != "complexity" + or (existing is not None and effective_model != existing_model) + ): + raise HTTPException(status_code=403, detail="Team members may manage only complexity auto routers.") + public_name: Final = ( + existing.model_info.team_public_model_name or existing.model_name + if existing is not None + else incoming.model_name + ) + if ( + not public_name + or public_name != public_name.strip() + or any(character in public_name for character in "*?[]") + or public_name.startswith("model_name_") + ): + raise HTTPException( + status_code=400, detail="Choose a non-empty auto-router name without wildcards or internal prefixes." + ) + if existing is not None and incoming.model_name not in (None, public_name, existing.model_name): + raise HTTPException(status_code=403, detail="Team members cannot rename an auto router.") + supplied_config: Final = _RouterConfigSource.model_validate(params.model_dump()).complexity_router_config + raw_config: Final = ( + supplied_config + if supplied_config is not None + else _RouterConfigSource.model_validate(existing.litellm_params.model_dump()).complexity_router_config + if existing is not None + else None + ) + if raw_config is None: + raise HTTPException(status_code=400, detail="A complexity_router_config is required.") + config: Final = validate_member_auto_router_config(raw_config) + stored_default: Final = existing.litellm_params.complexity_router_default_model if existing is not None else None + default_model: Final = ( + params.complexity_router_default_model + if params.complexity_router_default_model is not None + else decrypt_value_helper(stored_default, key="complexity_router_default_model", return_original_value=True) + if stored_default is not None + else None + ) + await authorize_member_auto_router_dependencies( + config=config, + default_model=default_model, + user_api_key_dict=user_api_key_dict, + team=team, + prisma_client=prisma_client, + llm_router=llm_router, + ) + return MemberAutoRouterWrite( + actor=user_api_key_dict, + team_id=team.team_id, + model_id=existing.model_info.id if existing is not None else None, + public_name=public_name, + updated_at=stored.updated_at if stored is not None else None, + config=config, + default_model=default_model, + ) diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index d962934dfb1..93b8c5c7cd7 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -12,6 +12,11 @@ from typing import Protocol, TypeVar RowT_co = TypeVar("RowT_co", covariant=True) +class DatabaseClient(Protocol): + @property + def db(self) -> object: ... + + class TableActions(Protocol[RowT_co]): """The prisma-client-py per-model action surface, keyed to the row it returns. diff --git a/litellm/router.py b/litellm/router.py index 0657c1e05ba..714a7f6248d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13463,7 +13463,7 @@ class Router: async def async_pre_routing_hook( self, model: str, - request_kwargs: dict, + request_kwargs: dict[str, object], messages: list[dict[str, Any]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, @@ -13511,6 +13511,18 @@ class Router: ) return None + from litellm.proxy.auth.auto_router_checks import authorize_member_auto_router_inference + + await authorize_member_auto_router_inference( + deployment=self._selected_strategy_marker_deployment( + model=registered_model_name, + strategy_tags=selected_strategy.tags, + request_kwargs=request_kwargs, + ), + request_kwargs=request_kwargs, + llm_router=self, + ) + from litellm.proxy.guardrails.auto_router_compression import ( messages_for_routing, model_hop_compression_armed, @@ -13610,25 +13622,34 @@ class Router: return pre_routing_hook_response + def _selected_strategy_marker_deployment( + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] + ) -> DeploymentTypedDict | None: + markers: Final = tuple( + deployment + for deployment in self.deployments_for_request(model, request_kwargs) + if "model" in deployment["litellm_params"] + and str(deployment["litellm_params"]["model"]).startswith(AUTO_ROUTER_MODEL_PREFIX) + ) + tag_matched: Final = tuple( + deployment + for deployment in markers + if (tuple(deployment["litellm_params"]["tags"] or ()) if "tags" in deployment["litellm_params"] else ()) + == strategy_tags + ) + return tag_matched[0] if tag_matched else (markers[0] if markers else None) + def _forwardable_alias_marker_params( self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] ) -> tuple[tuple[str, object], ...]: - marker_params: Final = tuple( - litellm_params - for deployment in self.deployments_for_request(model, request_kwargs) - if str((litellm_params := deployment["litellm_params"]).get("model", "")).startswith( - AUTO_ROUTER_MODEL_PREFIX - ) + marker: Final = self._selected_strategy_marker_deployment( + model=model, strategy_tags=strategy_tags, request_kwargs=request_kwargs ) - tag_matched: Final = tuple( - params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags - ) - selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) - if selected is None: + if marker is None: return () return tuple( (key, value) - for key, value in selected.items() + for key, value in marker["litellm_params"].items() if key not in _ALIAS_PARAMS_NEVER_FORWARDED and key not in CustomPricingLiteLLMParams.model_fields and value is not None diff --git a/litellm/types/router.py b/litellm/types/router.py index 0aefc07ae4b..edd42f264d1 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -180,6 +180,7 @@ class ModelInfo(MirroredPricingParams): # the model_name that can be used by the team when making LLM calls team_public_model_name: str | None = None + member_auto_router: bool = False # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked blocked: bool | None = None diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5f87f2def93..f5d5fc0f78b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5114,8 +5114,9 @@ async def test_model_discovery_route_bypasses_user_budget(): assert result is True +@pytest.mark.parametrize("route", ["/health/services", "/auto_router/test_routing"]) @pytest.mark.asyncio -async def test_side_effectful_info_route_still_enforces_budget(): +async def test_side_effectful_info_route_still_enforces_budget(route: str) -> None: """#27923 keeps the bypass narrow: /health/services can fire Slack/email/webhook test messages, so an exhausted budget must still block it. Widening the exemption back to is_info_route() would regress this.""" @@ -5131,7 +5132,7 @@ async def test_side_effectful_info_route_still_enforces_budget(): end_user_object=None, global_proxy_spend=None, general_settings={}, - route="/health/services", + route=route, llm_router=None, proxy_logging_obj=AsyncMock(), valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"), @@ -8146,3 +8147,33 @@ async def test_enforced_model_allowlists_reads_every_level_from_cache(): ] assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("channel", ["team", "key"]) +async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None: + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import can_key_call_model, can_team_access_model + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + group: Final = LiteLLM_AccessGroupTable( + access_group_id="group-a", access_group_name="allowed-models", access_model_names=["allowed"] + ) + reader: Final = AsyncMock(return_value=group) + client: Final = MagicMock(db=MagicMock(litellm_accessgrouptable=MagicMock(find_unique=reader))) + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: [TQ008] prove reads stay on the injected connection + patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), # test-quality-ok: [TQ008] isolate the process cache + ): + if channel == "team": + assert await can_team_access_model( + model="allowed", team_object=LiteLLM_TeamTable(team_id="team-a", models=["other"], access_group_ids=["group-a"]), + llm_router=None, prisma_client=client, + ) is True + else: + assert await can_key_call_model( + model="allowed", llm_model_list=None, + valid_token=UserAPIKeyAuth(models=["other"], access_group_ids=["group-a"]), + llm_router=None, prisma_client=client, + ) is True + reader.assert_awaited_once_with(where={"access_group_id": "group-a"}) diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index ef843adad98..fb7b515eca1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -7,7 +7,7 @@ from pathlib import Path from typing import Final import pytest -from fastapi import HTTPException +from fastapi import HTTPException, Request from pydantic import ValidationError from litellm.proxy._types import ( @@ -26,6 +26,8 @@ from litellm.types.management_endpoints.auto_router_endpoints import ( ) from litellm.types.utils import Choices, Message, ModelResponse +ROUTING_HTTP_REQUEST: Final = Request({"type": "http", "method": "POST", "path": "/auto_router/test_routing", "headers": []}) + ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test", user_id="admin") @@ -94,6 +96,7 @@ async def _route_body(body: Mapping[str, object], monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(proxy_server, "llm_router", _router()) return await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request_from(body, **config_overrides), user_api_key_dict=ADMIN, ) @@ -121,6 +124,7 @@ async def _classifier_user_payload(body: Mapping[str, object], monkeypatch: pyte monkeypatch.setattr(proxy_server, "llm_router", router) await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request_from(body, classifier_type="llm", classifier_llm_config={"model": "classifier-model"}), user_api_key_dict=ADMIN, ) @@ -198,6 +202,7 @@ async def test_llm_classifier_call_is_billed_to_the_calling_key(monkeypatch: pyt monkeypatch.setattr(proxy_server, "llm_router", router) response = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request( "what is 2+2", classifier_type="llm", @@ -359,6 +364,7 @@ async def test_a_key_that_cannot_call_the_classifier_model_is_rejected_before_it with pytest.raises(ProxyException) as exc_info: await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request("what is 2+2", **config_overrides), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, @@ -388,6 +394,7 @@ async def test_a_key_over_its_budget_cannot_run_a_classifier_config(monkeypatch: with pytest.raises(ProxyException) as exc_info: await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request( "what is 2+2", classifier_type="llm", @@ -413,6 +420,7 @@ async def test_a_heuristic_config_does_not_need_a_budget(monkeypatch: pytest.Mon monkeypatch.setattr(proxy_server, "llm_router", _router()) response = await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request("what is 2+2"), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, @@ -434,7 +442,7 @@ async def test_no_llm_router_on_the_proxy_is_a_500(monkeypatch: pytest.MonkeyPat monkeypatch.setattr(proxy_server, "llm_router", None) with pytest.raises(HTTPException) as exc_info: - await preview_auto_router_routing(data=_request("what is 2+2"), user_api_key_dict=ADMIN) + await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=_request("what is 2+2"), user_api_key_dict=ADMIN) assert exc_info.value.status_code == 500 @@ -447,6 +455,7 @@ async def test_non_admin_without_a_team_is_rejected(monkeypatch: pytest.MonkeyPa with pytest.raises(HTTPException) as exc_info: await preview_auto_router_routing( + http_request=ROUTING_HTTP_REQUEST, data=_request("what is 2+2"), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="user" @@ -2712,12 +2721,12 @@ async def test_routing_test_never_confirms_models_the_caller_cannot_use(monkeypa ) monkeypatch.setattr(proxy_server, "prisma_client", _team_prisma("team-probe", models=["mid-model"])) - probing = await preview_auto_router_routing(data=_request("team-probe"), user_api_key_dict=team_admin) + probing = await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=_request("team-probe"), user_api_key_dict=team_admin) assert probing.routed_model == "cheap-model" assert probing.routed_model_configured is False monkeypatch.setattr(proxy_server, "prisma_client", _team_prisma("team-grant", models=["cheap-model"])) - granted = await preview_auto_router_routing(data=_request("team-grant"), user_api_key_dict=team_admin) + granted = await preview_auto_router_routing(http_request=ROUTING_HTTP_REQUEST, data=_request("team-grant"), user_api_key_dict=team_admin) assert granted.routed_model == "cheap-model" assert granted.routed_model_configured is True @@ -2770,6 +2779,116 @@ async def test_validate_config_gates_like_the_write_it_rehearses(monkeypatch: py assert not_their_team.value.status_code == 403 +def _configure_member_preview( + monkeypatch: pytest.MonkeyPatch, *, allowed: bool = True +) -> UserAPIKeyAuth: + from litellm.proxy import proxy_server + from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTable + + team: Final = LiteLLM_TeamTable( + team_id="member-preview-team", + models=list(TIERS[name][0] for name in TIERS), + members_with_roles=[{"role": "user", "user_id": "preview-member"}], + team_member_permissions=["/auto_router/manage"] if allowed else [], + ) + prisma: Final = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + monkeypatch.setattr(proxy_server, "premium_user", True) + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="preview-member", + team_id=UI_TEAM_ID, + api_key="sk-preview-member", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("access", ["allowed", "opt-out", "limited-key"]) +async def test_member_preview_and_validation_follow_team_opt_in( + monkeypatch: pytest.MonkeyPatch, access: str +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import validate_complexity_router_config + from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest + + actor: Final = _configure_member_preview(monkeypatch, allowed=access != "opt-out").model_copy(update={ + "models": ["member-router"] if access == "limited-key" else [], "config": {"timeout": 60}, + }) + monkeypatch.setattr(proxy_server, "llm_router", _router()) + preview: Final = _request_from({"prompt": "what is 2+2", "team_id": "member-preview-team"}) + validation: Final = ComplexityRouterConfigValidationRequest( + team_id="member-preview-team", complexity_router_config={"tiers": TIERS, "classifier_type": "heuristic"} + ) + if access != "allowed": + with pytest.raises((HTTPException, ProxyException)) as denied_preview: + await preview_auto_router_routing(preview, actor, ROUTING_HTTP_REQUEST) + with pytest.raises((HTTPException, ProxyException)) as denied_validation: + await validate_complexity_router_config(validation, actor) + assert str(getattr(denied_preview.value, "status_code", None) or denied_preview.value.code) == "403" + assert str(getattr(denied_validation.value, "status_code", None) or denied_validation.value.code) == "403" + return + assert (await validate_complexity_router_config(validation, actor)).valid is True + result: Final = await preview_auto_router_routing(preview, actor, ROUTING_HTTP_REQUEST) + assert result.routed_model == "cheap-model" + assert result.routed_model_configured is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("over_budget", [False, True]) +async def test_member_billable_preview_checks_and_charges_destination_team( + monkeypatch: pytest.MonkeyPatch, over_budget: bool +) -> None: + import importlib + + import litellm + from litellm.proxy import proxy_server + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + auth_module: Final = importlib.import_module("litellm.proxy.auth.user_api_key_auth") + actor: Final = _configure_member_preview(monkeypatch).model_copy(update={"metadata": {"tags": ["key-tag"]}}) + router: Final = RecordingRouter("SIMPLE") + monkeypatch.setattr(proxy_server, "llm_router", router) + + async def check_and_tag( + user_api_key_auth_obj: UserAPIKeyAuth, request: Request, request_data: dict[str, object], route: str + ) -> None: + assert route == "/auto_router/test_routing" + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request, request_data=request_data, user_api_key_dict=user_api_key_auth_obj + ) + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( + request_data=request_data, user_api_key_dict=user_api_key_auth_obj + ) + if over_budget: + raise litellm.BudgetExceededError(current_cost=2, max_budget=1) + + checks: Final = AsyncMock(side_effect=check_and_tag) + monkeypatch.setattr(auth_module, "_run_centralized_common_checks", checks) + http_request: Final = Request({ + "type": "http", "method": "POST", "path": "/auto_router/test_routing", + "headers": [(b"x-litellm-tags", b"header-tag")], + }) + data: Final = _request_from( + {"prompt": "hi", "team_id": "member-preview-team"}, + classifier_type="llm", classifier_llm_config={"model": "cheap-model"}, + ) + if over_budget: + with pytest.raises(litellm.BudgetExceededError): + await preview_auto_router_routing(data, actor, http_request) + assert router.recorded_calls == [] + else: + await preview_auto_router_routing(data, actor, http_request) + assert len(router.recorded_calls) == 1 + assert router.recorded_calls[0]["metadata"]["user_api_key_team_id"] == "member-preview-team" + assert router.recorded_calls[0]["metadata"]["user_api_key_user_id"] == "preview-member" + assert set(router.recorded_calls[0]["metadata"]["tags"]) == {"key-tag", "header-tag"} + checks.assert_awaited_once() + assert checks.await_args.kwargs["user_api_key_auth_obj"].team_id == "member-preview-team" + assert checks.await_args.kwargs["route"] == "/auto_router/test_routing" + + def test_every_shadow_eval_sql_constant_speaks_naive_utc(): """The tables store naive UTC wall time (prisma's convention), so SQL-side time must be NOW() AT TIME ZONE 'utc' and python-side params must cast ::timestamp; a bare NOW() or a diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index c3ad66397ea..6300331d564 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -2,7 +2,7 @@ import inspect import asyncio import contextlib import json -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from typing import Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -1404,7 +1404,7 @@ class TestTeamModelSiblingRouting: side_effect=mock_add_model_to_db, ), patch( - "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + "litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", mock_team_model_add, ), ): @@ -5323,7 +5323,7 @@ class TestStrategyRouterWriteValidation: lambda value, new_encryption_key=None: value, ), patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted - "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + "litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", side_effect=team_model_add, ), ): @@ -6198,3 +6198,232 @@ class TestAccessGroupModelSync: assert "array_replace" in update_call.args[0] assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") invalidate.assert_awaited_once_with(("ag-1",)) + + +class TestTeamMemberAutoRouterWrites: + @pytest.fixture(autouse=True) + def _salt(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "member-router-test-salt") + + @contextlib.contextmanager + def _environment(self, database: MagicMock, row: LiteLLM_ProxyModelTable) -> Iterator[None]: + with ( + patch("litellm.proxy.proxy_server.prisma_client", database), # test-quality-ok: [TQ008] endpoint storage singleton injection + patch("litellm.proxy.proxy_server.llm_router", self._catalog()), # test-quality-ok: [TQ008] inject real destination model catalog + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint storage mode singleton + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] inject licensed process state + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", return_value=None), # test-quality-ok: [TQ008] inject unlimited license result + patch("litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", new=AsyncMock()), # test-quality-ok: [TQ008] pubsub I/O boundary + patch("litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", new=AsyncMock()), # test-quality-ok: [TQ008] audit database I/O boundary + patch("litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock(return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary + still_desired=frozenset((row.model_id, "allowed-id")), live_after=frozenset((row.model_id, "allowed-id")) + ))), + ): + yield + + @staticmethod + def _team(enabled: bool = True) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable( + team_id="member-team", + models=["allowed"], + members_with_roles=[Member(user_id="owner", role="user"), Member(user_id="peer", role="user")], + team_member_permissions=["/auto_router/manage"] if enabled else [], + ) + + @staticmethod + def _row() -> LiteLLM_ProxyModelTable: + return LiteLLM_ProxyModelTable( + model_id="member-router", + model_name="model_name_member-team_stored", + litellm_params={ + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "allowed"}}, + "complexity_router_default_model": "allowed", + }, + model_info={ + "id": "member-router", + "team_id": "member-team", + "team_public_model_name": "personal-router", + "created_by": "peer", + "access_groups": ["retained-admin-group"], + }, + created_by="owner", + ) + + @staticmethod + def _database(team: LiteLLM_TeamTable, row: LiteLLM_ProxyModelTable) -> MagicMock: + table: Final = MagicMock( + find_unique=AsyncMock(return_value=row), + find_many=AsyncMock(return_value=[]), + update=AsyncMock(return_value=row), + create=AsyncMock(return_value=row), + ) + transaction: Final = MagicMock( + litellm_teamtable=MagicMock(find_unique=AsyncMock(return_value=team)), + litellm_teammembership=MagicMock(find_unique=AsyncMock(return_value=None)), + litellm_proxymodeltable=table, + query_raw=AsyncMock(return_value=[]), + ) + context: Final = MagicMock( + __aenter__=AsyncMock(return_value=transaction), + __aexit__=AsyncMock(return_value=False), + ) + db: Final = MagicMock( + litellm_teamtable=MagicMock(find_unique=AsyncMock(return_value=team)), + litellm_teammembership=MagicMock(find_unique=AsyncMock(return_value=None)), + litellm_proxymodeltable=table, + tx=MagicMock(return_value=context), + ) + return MagicMock(db=db, transaction=transaction) + + @staticmethod + def _catalog() -> Router: + return Router(model_list=[{ + "model_name": "allowed", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, + "model_info": {"id": "allowed-id"}, + }]) + + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint,change", [("patch", "config"), ("legacy", "strategy"), ("patch", "unrelated")]) + async def test_admin_router_changes_release_member_scope(self, endpoint: str, change: str) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model, update_model + + original: Final = self._row() + row: Final = original.model_copy(update={"model_info": {**original.model_info, "member_auto_router": True}}) + database: Final = self._database(self._team(), row) + params: Final = { + "config": {"complexity_router_config": {"tiers": {"SIMPLE": "allowed"}, "session_affinity": True}}, + "strategy": {"model": "auto_router/quality_router", "quality_router_default_model": "allowed"}, + "unrelated": {"model": "auto_router/complexity_router", "max_tokens": 100}, + } + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams.model_validate(params[change]), + model_info=ModelInfo(id=row.model_id) if endpoint == "legacy" or change == "unrelated" else None, + ) + with self._environment(database, row): + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + if endpoint == "patch": + await patch_model(row.model_id, request, actor) + else: + await update_model(request, actor) + written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"] + saved_info: Final = json.loads(written["model_info"]) if "model_info" in written else row.model_info + assert saved_info["member_auto_router"] is (change == "unrelated") + assert saved_info["team_id"] == "member-team" + assert saved_info["access_groups"] == ["retained-admin-group"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) + @pytest.mark.parametrize("access", ["owner", "peer", "limited-key"]) + async def test_both_update_entries_enforce_creator_and_stamp_member_scope( + self, endpoint: str, access: str + ) -> None: + from fastapi import HTTPException + + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model, update_model + + row: Final = self._row() + database: Final = self._database(self._team(), row) + request: Final = updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config={"tiers": {"SIMPLE": "allowed"}, "session_affinity": True}), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ) + actor: Final = UserAPIKeyAuth( + user_id="peer" if access == "peer" else "owner", user_role=LitellmUserRoles.INTERNAL_USER, + models=["personal-router"] if access == "limited-key" else ["allowed"], config={"timeout": 60}, + ) + with self._environment(database, row): + operation: Final = patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + if access != "owner": + with pytest.raises((HTTPException, ProxyException)): + await operation + database.transaction.litellm_proxymodeltable.update.assert_not_awaited() + return + await operation + written: Final = database.transaction.litellm_proxymodeltable.update.await_args.kwargs["data"] + saved_info: Final = json.loads(written["model_info"]) + assert saved_info["member_auto_router"] is True + assert saved_info["team_id"] == "member-team" + assert saved_info["access_groups"] == ["retained-admin-group"] + assert "created_by" not in written + assert json.loads(written["litellm_params"])["complexity_router_config"]["session_affinity"] is True + assert written.get("model_name", row.model_name) == row.model_name + + @pytest.mark.asyncio + @pytest.mark.parametrize("changed_state", ["allowed", "revoked", "moved", "creator", "collision", "global-alias"]) + async def test_write_slot_rechecks_authoritative_team_owner_and_names(self, changed_state: str) -> None: + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + from litellm.proxy.management_helpers.auto_router_permissions import MemberAutoRouterWrite, validate_member_auto_router_config + + row: Final = self._row() + database: Final = self._database(self._team(), row) + if changed_state == "revoked": + database.transaction.litellm_teamtable.find_unique.return_value = self._team(enabled=False) + elif changed_state == "moved": + database.transaction.litellm_proxymodeltable.find_unique.return_value = row.model_copy(update={"model_info": {"team_id": "other-team"}}) + elif changed_state == "creator": + database.transaction.litellm_proxymodeltable.find_unique.return_value = row.model_copy(update={"created_by": "peer"}) + elif changed_state == "collision": + database.transaction.litellm_proxymodeltable.find_many.return_value = [row] + config: Final = validate_member_auto_router_config({"tiers": {"SIMPLE": "allowed"}}) + grant: Final = MemberAutoRouterWrite( + actor=UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, models=["allowed"]), + team_id="member-team", model_id=None if changed_state in ("collision", "global-alias") else row.model_id, + public_name="personal-router", updated_at=None, config=config, default_model="allowed", + ) + with ( + self._environment(database, row), + patch("litellm.model_alias_map", {"personal-router": "allowed"} if changed_state == "global-alias" else {}), # test-quality-ok: [TQ008] inject alias namespace for collision behavior + ): + if changed_state != "allowed": + with pytest.raises(HTTPException) as denied: + async with _auto_router_capability_slot(database, effective_params={}, model_id=grant.model_id, member_write=grant): + pytest.fail("An invalidated grant reached the database writer") + assert denied.value.status_code == (409 if changed_state in ("collision", "global-alias") else 403) + return + async with _auto_router_capability_slot(database, effective_params={}, model_id=grant.model_id, member_write=grant) as table: + await table.update(where={"model_id": row.model_id}, data={"updated_by": "owner"}) + assert database.transaction.query_raw.await_count == 2 + database.transaction.litellm_proxymodeltable.update.assert_awaited_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("access", ["allowed", "opt-out", "limited-key"]) + async def test_create_entry_requires_opt_in_and_appends_only_its_router(self, access: str) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + row: Final = self._row() + database: Final = self._database(self._team(enabled=access != "opt-out"), row) + actor: Final = UserAPIKeyAuth( + user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, + models=["personal-router"] if access == "limited-key" else ["allowed"], config={"timeout": 60}, + ) + deployment: Final = Deployment( + model_name="new-personal-router", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config={"tiers": {"SIMPLE": "allowed"}}), + model_info=ModelInfo(id=row.model_id, team_id="member-team"), + ) + with ( + self._environment(database, row), + patch("litellm.proxy.proxy_server.proxy_config.add_deployment", new=AsyncMock(return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary + still_desired=frozenset((row.model_id, "allowed-id")), live_after=frozenset((row.model_id, "allowed-id")) + ))), + patch("litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", new=AsyncMock()) as appended, # test-quality-ok: [TQ008] persistence boundary; the appended scope is asserted + ): + if access != "allowed": + with pytest.raises(ProxyException) as denied: + await add_new_model(deployment, actor) + assert denied.value.code == "403" + database.transaction.litellm_proxymodeltable.create.assert_not_awaited() + appended.assert_not_awaited() + return + await add_new_model(deployment, actor) + written: Final = database.transaction.litellm_proxymodeltable.create.await_args.kwargs["data"] + assert written["created_by"] == "owner" + assert json.loads(written["model_info"])["member_auto_router"] is True + assert appended.await_args.kwargs["data"].models == ["new-personal-router"] + assert appended.await_args.kwargs["data"].team_id == "member-team" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 9a4badab8a9..2e0ef52ed38 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -8914,6 +8914,11 @@ async def test_delete_team_survives_a_failing_cache_backend( @pytest.mark.asyncio async def test_team_member_delete_persists_deleted_keys(monkeypatch): from litellm.proxy._types import TeamMemberDeleteRequest + from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, + ) from litellm.proxy.management_endpoints.key_management_endpoints import ( LiteLLM_VerificationToken, ) @@ -9011,6 +9016,16 @@ async def test_team_member_delete_persists_deleted_keys(monkeypatch): lambda **kwargs: True, ) + cache: Final = UserApiKeyCache() + revoked_cache_keys: Final = ( + "team_id:team-1", "team_alias:test-team", "user-123", "hashed-token-1", "hashed-token-2", + team_membership_auth_cache_key(user_id="user-123", team_id="team-1"), + team_membership_reservation_cache_key(user_id="user-123", team_id="team-1"), + ) + for cache_key in (*revoked_cache_keys, "unrelated-key"): + cache.set_cache(key=cache_key, value={"retained": True}) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + data = TeamMemberDeleteRequest(team_id="team-1", user_id="user-123") result = await team_member_delete( @@ -9027,6 +9042,9 @@ async def test_team_member_delete_persists_deleted_keys(monkeypatch): assert all(record["team_id"] == "team-1" for record in records) assert all(record["user_id"] == "user-123" for record in records) mock_delete_keys.assert_called_once() + assert result.members_with_roles == [] + assert all(cache.get_cache(key=cache_key) is None for cache_key in revoked_cache_keys) + assert cache.get_cache(key="unrelated-key") == {"retained": True} @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_helpers/test_auto_router_permissions.py b/tests/test_litellm/proxy/management_helpers/test_auto_router_permissions.py new file mode 100644 index 00000000000..fb91a23088c --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_auto_router_permissions.py @@ -0,0 +1,208 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ( + UI_TEAM_ID, + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + UserAPIKeyAuth, +) +from litellm.proxy.management_helpers.auto_router_permissions import ( + authorize_member_auto_router_dependencies, + authorize_member_auto_router_team, + authorize_member_auto_router_write, + validate_member_auto_router_config, +) +from litellm.router import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment + + +class _ReadTable: + async def find_unique( + self, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> None: + return None + + +@dataclass(frozen=True) +class _PermissionDb: + litellm_teammembership: _ReadTable = _ReadTable() + + +@dataclass(frozen=True) +class _Client: + db: _PermissionDb = _PermissionDb() + + +def _team(**updates: object) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable.model_validate( + { + "team_id": "team-a", + "models": ["allowed"], + "members_with_roles": [Member(user_id="owner", role="user")], + "team_member_permissions": ["/auto_router/manage"], + **updates, + } + ) + + +def _actor(**updates: object) -> UserAPIKeyAuth: + return UserAPIKeyAuth.model_validate( + {"user_id": "owner", "user_role": "internal_user", "models": ["allowed"], **updates} + ) + + +@pytest.fixture +def catalog() -> Router: + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}} + for name in ("allowed", "other") + ] + ) + + +@pytest.mark.parametrize( + "actor_updates,team_updates,premium,allowed", + [ + ({}, {}, True, True), + ({"team_id": UI_TEAM_ID}, {}, True, True), + ({"team_id": "team-a"}, {}, True, True), + ({"user_role": LitellmUserRoles.TEAM}, {}, True, True), + ({"user_role": LitellmUserRoles.ORG_ADMIN}, {}, True, True), + ({"team_id": "team-b"}, {}, True, False), + ({"user_id": None}, {}, True, False), + ({"user_id": ""}, {}, True, False), + ({"user_id": "peer"}, {}, True, False), + ({"user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY}, {}, True, False), + ({"user_role": LitellmUserRoles.CUSTOMER}, {}, True, False), + ({}, {"team_member_permissions": []}, True, False), + ({}, {"team_member_permissions": None}, True, False), + ({}, {"blocked": True}, True, False), + ({}, {}, False, False), + ], +) +def test_opt_in_requires_live_named_membership_and_write_role( + actor_updates: Mapping[str, object], team_updates: Mapping[str, object], premium: bool, allowed: bool +) -> None: + if allowed: + authorize_member_auto_router_team( + user_api_key_dict=_actor(**actor_updates), team=_team(**team_updates), premium_user=premium + ) + return + with pytest.raises(HTTPException) as denied: + authorize_member_auto_router_team( + user_api_key_dict=_actor(**actor_updates), team=_team(**team_updates), premium_user=premium + ) + assert denied.value.status_code == 403 + + +@pytest.mark.parametrize("placement", ["inline", "normalized"]) +@pytest.mark.parametrize( + "overrides", [{"api_base": "https://example.invalid"}, {"api_key": "fake"}, {"metadata": {}}, {"model": "other"}] +) +def test_all_tier_parameter_representations_reject_privileged_overrides( + placement: str, overrides: Mapping[str, object] +) -> None: + entry: Final = {"model_name": "allowed", "litellm_params": overrides} + config: Final = ( + {"tiers": {"SIMPLE": [entry]}} + if placement == "inline" + else {"tiers": {"SIMPLE": ["allowed"]}, "tier_model_configs": {"SIMPLE": [entry]}} + ) + with pytest.raises(HTTPException) as denied: + validate_member_auto_router_config(config) + assert denied.value.status_code == 400 + + +def test_tier_config_is_normalized_and_unknown_router_extras_are_rejected() -> None: + validated: Final = validate_member_auto_router_config( + {"tiers": {"SIMPLE": [{"model_name": "allowed", "litellm_params": {"reasoning_effort": "low"}}]}} + ) + assert validated.tiers == {"SIMPLE": ["allowed"]} + assert validated.tier_model_configs["SIMPLE"][0].litellm_params == {"reasoning_effort": "low"} + assert validate_member_auto_router_config(validated.model_dump()).tiers == validated.tiers + with pytest.raises(HTTPException): + validate_member_auto_router_config({"tiers": {"SIMPLE": "allowed"}, "api_base": "https://example.invalid"}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "patch_fields", + [ + {}, + {"model_name": "renamed"}, + {"blocked": False}, + {"model_info": {"team_id": "other-team"}}, + {"model_info": {"member_auto_router": False}}, + {"litellm_params": {"model": "auto_router/quality_router"}}, + {"litellm_params": {"api_key": "fake"}}, + ], +) +async def test_member_updates_restrict_fields_and_preserve_an_inherited_default( + catalog: Router, monkeypatch: pytest.MonkeyPatch, patch_fields: Mapping[str, object] +) -> None: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + monkeypatch.setenv("LITELLM_SALT_KEY", "member-router-test-salt") + existing: Final = Deployment( + model_name="model_name_team-a_uuid", + litellm_params=LiteLLM_Params( + model=encrypt_value_helper("auto_router/complexity_router"), + complexity_router_config={"tiers": {"SIMPLE": "allowed"}}, + complexity_router_default_model=encrypt_value_helper("allowed"), + ), + model_info=ModelInfo(id="router-a", team_id="team-a", team_public_model_name="my-router"), + created_by="owner", + ) + patch: Final = updateDeployment.model_validate( + {"litellm_params": {"complexity_router_config": {"tiers": {"SIMPLE": "allowed"}}}, **patch_fields} + ) + operation: Final = authorize_member_auto_router_write( + incoming=patch, + existing=existing, + user_api_key_dict=_actor(), + team=_team(), + premium_user=True, + prisma_client=_Client(), + llm_router=catalog, + ) + if patch_fields: + with pytest.raises(HTTPException) as denied: + await operation + assert denied.value.status_code == 403 + return + granted: Final = await operation + assert granted.default_model == "allowed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target", ["missing", "nested"]) +async def test_member_dependencies_require_plain_configured_models(target: str) -> None: + catalog: Final = Router( + model_list=[ + {"model_name": "allowed", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}}, + { + "model_name": "nested", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "allowed"}}, + }, + }, + ] + ) + with pytest.raises(HTTPException) as denied: + await authorize_member_auto_router_dependencies( + config=validate_member_auto_router_config({"tiers": {"SIMPLE": target}}), + default_model=None, + user_api_key_dict=_actor(models=[target]), + team=_team(models=[target]), + prisma_client=_Client(), + llm_router=catalog, + ) + assert denied.value.status_code == 400 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f1b445fb1bd..40787ead91f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4,9 +4,10 @@ import functools import json import logging import os +import sys import threading -from datetime import datetime from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -15,36 +16,37 @@ import httpx import openai import pytest import respx - - +from fastapi import HTTPException import litellm +from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.redis_cache import _redis_circuit_breaker_guard -from litellm import Router from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, ) -from litellm.types.llms.openai import ChatCompletionRequest +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, ProxyException, UserAPIKeyAuth from litellm.router import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, FallbackAwareAnthropicMessagesStream, _anthropic_stream_commits_now, + _anthropic_stream_error_is_gateway_verdict, _anthropic_stream_fallback_error_for_raised, + _anthropic_stream_forwards_ping_live, _anthropic_stream_raised_error_status, _anthropic_stream_should_decline_fallback, - _anthropic_stream_error_is_gateway_verdict, - _anthropic_stream_forwards_ping_live, _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, ) from litellm.router_strategy import simple_shuffle -from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, RetryPolicy +from litellm.types.llms.openai import ChatCompletionRequest +from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -15806,3 +15808,234 @@ async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warni assert binding is None assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) + + +class TestMemberAutoRouterInference: + @pytest.fixture(autouse=True) + def runtime(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + self.cache = UserApiKeyCache() + self.team = LiteLLM_TeamTable( + team_id="router-team", models=["member-router", "permitted-model"], + members_with_roles=[Member(user_id="router-member", role="user")], + ) + self.actor = UserAPIKeyAuth( + user_id="router-member", team_id="router-team", user_role=LitellmUserRoles.INTERNAL_USER, + models=["member-router", "permitted-model"], api_key="test-key-hash", config={"timeout": 60}, + ) + self.database = SimpleNamespace(db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), + litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), + )) + monkeypatch.setattr(proxy_server, "user_api_key_cache", self.cache) + monkeypatch.setattr(proxy_server, "prisma_client", self.database) + + @staticmethod + def _marker(*, member: bool = True, classifier: bool = False) -> dict[str, object]: + target: Final = "permitted-model" if member else "restricted-model" + return { + "model_name": "model_name_router-team_member-router", + "litellm_params": { + "model": "auto_router/complexity_router", "complexity_router_default_model": target, + "complexity_router_config": { + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), "adaptive": False, + **({"classifier_type": "llm", "classifier_llm_config": {"model": target}} if classifier else {}), + }, + "tags": ["member" if member else "admin"], "timeout": 13.0 if member else 29.0, + }, + "model_info": { + "team_id": "router-team", "team_public_model_name": "member-router", "member_auto_router": member, + }, + } + + @classmethod + def _router(cls, *markers: dict[str, object]) -> Router: + return Router(model_list=[ + *(markers or (cls._marker(),)), + {"model_name": "permitted-model", "litellm_params": { + "model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://api.openai.com/v1", + }}, + {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, + ]) + + def _request( + self, *, actor: UserAPIKeyAuth | None = None, metadata_name: str = "metadata", tag: str = "member", + ) -> dict[str, object]: + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={metadata_name: {"tags": [tag]}, **({"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} + if metadata_name == "litellm_metadata" else {})}, + user_api_key_dict=actor or self.actor, _metadata_variable_name=metadata_name, + ) + + async def _route( + self, router: Router, request: dict[str, object] | None = None, model: str = "member-router", + ) -> PreRoutingHookResponse: + response: Final = await router.async_pre_routing_hook( + model=model, request_kwargs=request if request is not None else self._request(), + messages=[{"role": "user", "content": "Hello"}], + ) + assert response is not None + return response + + @pytest.mark.asyncio + @pytest.mark.parametrize("metadata_name", ("metadata", "litellm_metadata")) + async def test_cached_roster_revocation_blocks_classifier_and_session_rebinding( + self, metadata_name: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + ) -> None: + from litellm.proxy.auth.auth_checks import delete_cache_team_object + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + router: Final = self._router(self._marker(classifier=True)) + classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond(200, json={ + "id": "classifier", "object": "chat.completion", "created": 0, "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) + request: Final = {**self._request(metadata_name=metadata_name), "proxy_server_request": {"headers": { + "x-claude-code-session-id": "member-router-session", "x-app": "cli", + }}} + first: Final = await self._route(router, request) + assert first.model == "permitted-model" and first.routing_decision is not None + assert first.routing_decision["cause"] == "llm_classifier" + assert (await self._route(router, request)).model == "permitted-model" + assert self.database.db.litellm_teamtable.find_unique.await_count == 1 + assert self.database.db.litellm_teammembership.find_unique.await_count == 1 + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={"members_with_roles": []}) + await delete_cache_team_object( + team_id=self.team.team_id, team_alias=None, user_api_key_cache=self.cache, proxy_logging_obj=None, + ) + with pytest.raises(HTTPException, match="no longer a member"): + await self._route(router, request) + rebound: Final = {**request, "proxy_server_request": {"headers": { + "x-claude-code-session-id": "member-router-session", "x-app": "cli", "x-claude-code-agent-id": "subagent", + }}} + with pytest.raises(HTTPException, match="no longer a member"): + await self._route(router, rebound, model="restricted-model") + assert classify.call_count == 2 + + @pytest.mark.asyncio + @pytest.mark.parametrize("state", ("forged", "blocked", "deleted", "unavailable", "empty-user")) + async def test_member_router_fails_closed(self, state: str, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + request: Final = {"metadata": {"user_api_key_team_id": "router-team", "user_api_key_auth": { + "team_id": "router-team", "user_role": "proxy_admin", + }}} if state == "forged" else self._request(actor=self.actor.model_copy( + update={"user_id": ""} if state == "empty-user" else {}, + )) + self.database.db.litellm_teamtable.find_unique.return_value = ( + None if state == "deleted" else self.team.model_copy(update={"blocked": state == "blocked"}) + ) + if state == "unavailable": + monkeypatch.setattr(proxy_server, "prisma_client", None) + with pytest.raises(HTTPException) as error: + await self._route(self._router(), request) + assert error.value.status_code == (503 if state == "unavailable" else 403) + + @pytest.mark.asyncio + @pytest.mark.parametrize("user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)]) + async def test_service_key_and_admin_preserve_runtime_access(self, user_id: str | None, role: LitellmUserRoles) -> None: + assert (await self._route(self._router(), self._request( + actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), + ))).model == "permitted-model" + + @pytest.mark.asyncio + @pytest.mark.parametrize("ceiling", ("team", "key", "member", "organization", "project")) + async def test_runtime_dependency_ceilings_use_cached_auth_state(self, ceiling: str) -> None: + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.organization import LiteLLM_OrganizationTable + from litellm.models.team_membership import LiteLLM_TeamMembership + from litellm.proxy._types import LiteLLM_ProjectTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ + "models": ["member-router"] if ceiling == "team" else self.team.models, + "organization_id": "router-org" if ceiling == "organization" else None, + }) + if ceiling == "member": + await self.cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="router-member", team_id="router-team"), + value=LiteLLM_TeamMembership(user_id="router-member", team_id="router-team", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"])), + model_type=LiteLLM_TeamMembership, + ) + elif ceiling == "organization": + await self.cache.async_set_cache( + key="org_id:router-org", value=LiteLLM_OrganizationTable( + organization_id="router-org", budget_id="org-budget", created_by="admin", updated_by="admin", + models=["restricted-model"], + ), model_type=LiteLLM_OrganizationTable, + ) + elif ceiling == "project": + await self.cache.async_set_cache( + key="project_id:router-project", value=LiteLLM_ProjectTableCachedObj( + project_id="router-project", team_id="router-team", models=["restricted-model"], + ), model_type=LiteLLM_ProjectTableCachedObj, + ) + with pytest.raises(ProxyException, match="not allowed to access model"): + await self._route(self._router(), self._request(actor=self.actor.model_copy(update={ + "models": ["member-router"] if ceiling == "key" else self.actor.models, + "project_id": "router-project" if ceiling == "project" else None, + }))) + + @pytest.mark.asyncio + @pytest.mark.parametrize("group_owner", ("team", "key")) + async def test_access_group_grants_are_cached_and_revoked(self, group_owner: str) -> None: + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast + + group: Final = LiteLLM_AccessGroupTable( + access_group_id="router-group", access_group_name="Router targets", access_model_names=["permitted-model"], + ) + self.database.db.litellm_accessgrouptable.find_unique.return_value = group + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ + "models": ["member-router"] if group_owner == "team" else self.team.models, + "access_group_ids": ["router-group"] if group_owner == "team" else [], + }) + request: Final = self._request(actor=self.actor.model_copy(update={ + "models": ["member-router"] if group_owner == "key" else self.actor.models, + "access_group_ids": ["router-group"] if group_owner == "key" else [], + })) + router: Final = self._router() + assert (await self._route(router, request)).model == "permitted-model" + assert (await self._route(router, request)).model == "permitted-model" + assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 1 + self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy(update={"access_model_names": []}) + await evict_and_broadcast(cache_keys=("access_group_id:router-group",), user_api_key_cache=self.cache) + with pytest.raises(ProxyException, match="not allowed to access model"): + await self._route(router, request) + assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 2 + + @pytest.mark.asyncio + async def test_tagged_marker_owns_authorization_and_forwarded_parameters(self) -> None: + router: Final = self._router(self._marker(member=False), self._marker()) + request: Final = self._request() + selected: Final = router._selected_strategy_marker_deployment( + model="model_name_router-team_member-router", strategy_tags=("member",), request_kwargs=request, + ) + assert selected is not None and selected["model_info"]["member_auto_router"] is True + assert (await self._route(router, request)).model == "permitted-model" + assert request["timeout"] == 13.0 + await self.cache.async_set_cache( + key="team_id:router-team", model_type=LiteLLM_TeamTable, + value=self.team.model_copy(update={"models": ["member-router"]}), + ) + with pytest.raises(ProxyException, match="not allowed to access model"): + await self._route(router, self._request()) + self.database.db.litellm_teamtable.find_unique.reset_mock() + admin: Final = self._request(tag="admin") + assert (await self._route(router, admin)).model == "restricted-model" + assert admin["timeout"] == 29.0 + self.database.db.litellm_teamtable.find_unique.assert_not_awaited() + + @pytest.mark.asyncio + async def test_sdk_router_does_not_import_proxy_dependencies(self, monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = self._router(self._marker(member=False)) + monkeypatch.setitem(sys.modules, "fastapi", None) + monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False) + assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model" diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index bd38eecd1c6..f097e6f58e5 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -1379,7 +1379,7 @@ def test_a_discarded_router_stops_contributing_to_later_reloads(monkeypatch): _invalidate_model_cost_lowercase_map() -def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered(): +def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: """ The rebuild is only correct if it reproduces the entries the original registration wrote, including the pieces that are derived rather than stored: @@ -1406,6 +1406,7 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered(): at_boot = copy.deepcopy(litellm.model_cost["priced-id"]) assert at_boot["input_cost_per_token"] == 0.000123 assert at_boot["cache_read_input_token_cost"] is not None + assert "member_auto_router" not in litellm.model_cost["gpt-4o"] _simulate_price_data_reload( copy.deepcopy(fetched_catalog), @@ -1416,9 +1417,11 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered(): f"the rebuild changed or dropped a field the boot registration wrote: " f"{ {k: (v, rebuilt.get(k)) for k, v in at_boot.items() if rebuilt.get(k) != v} }" ) - # The rebuild goes through the deployment stored in model_list, which also - # carries the router's own db_model flag; add_deployment already registers it. - assert set(rebuilt) - set(at_boot) <= {"db_model"} + assert {field: rebuilt[field] for field in set(rebuilt) - set(at_boot)} == { + "db_model": False, + "member_auto_router": False, + } + assert "member_auto_router" not in litellm.model_cost["gpt-4o"] assert router.model_list finally: litellm.model_cost = saved_catalog diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx index 5b53217f9c1..f11b74c939d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx @@ -114,6 +114,7 @@ export function AutoRoutersPanel({ userRole={userRole} userId={userID} createScope={createScope} + teams={teams} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts index c4d7f45b7cc..7821437441d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.ts @@ -7,7 +7,7 @@ import { } from "@/components/add_model/auto_router_strategies"; import { normalizeTierModels } from "@/components/add_model/complexity_router_tiers"; import { Team } from "@/components/networking"; -import { type ModelActor, canModifyModel } from "@/utils/modelPermissions"; +import { type ModelActor, canEditAutoRouter, canModifyModel } from "@/utils/modelPermissions"; export type { AutoRouterKind }; @@ -106,13 +106,20 @@ export const toAutoRouterRow = ( const name = deployment.model_name ?? ""; const strategy = autoRouterStrategy(params); const { canEdit, canDelete, editBlockedReason } = autoRouterCapabilities(params, info); - const mayActOnRow = canModifyModel(actor, teams, { teamId: info.team_id, isDbModel: info.db_model === true }); + const origin = { + teamId: info.team_id, + isDbModel: info.db_model === true, + createdBy: info.created_by, + model: params.model, + }; + const mayActOnRow = canModifyModel(actor, teams, origin); + const mayEditRouter = canEditAutoRouter(actor, teams, origin); return { id: info.id ?? `${name}-${index}`, name, kind: strategy.kind, - canEdit: canEdit && mayActOnRow, + canEdit: canEdit && mayEditRouter, canDelete: canDelete && mayActOnRow, editBlockedReason, createdAt: info.created_at ?? undefined, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 4d6a90fc56e..bbc803af700 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -7,7 +7,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { all_admin_roles, internalUserRoles } from "@/utils/roles"; -import { canCreateModels } from "@/utils/modelPermissions"; +import { autoRouterCreationScope, canCreateModels } from "@/utils/modelPermissions"; import BetaBadge from "@/components/BetaBadge"; import CostOptimizationFeedbackBanner from "@/components/molecules/cost_optimization_feedback_banner"; import ModelInfoView from "@/components/model_info_view"; @@ -100,12 +100,17 @@ export default function ModelsAndEndpointsPage() { }, ); const isAdmin = all_admin_roles.includes(userRole); + const canViewAutoRouters = + autoRouterCreationScope( + { userRole, userID, isViewOnly }, + { teams: teams ?? null, disabledForInternalUsers: false }, + ) !== "forbidden"; const visibleSlugs = useMemo>( () => [ "", ...(canCreate ? (["add"] as const) : []), - ...(isAdmin || canCreate ? (["auto-routers"] as const) : []), + ...(isAdmin || canViewAutoRouters ? (["auto-routers"] as const) : []), // effectiveSessionRole reports proxy_admin_viewer as "Admin", so isAdmin alone would show a // viewer these write-only panels; only the raw-role isViewOnly separates them. Health Status // stays: it is the bucket's one read view, and viewers keep read parity with admins. @@ -115,7 +120,7 @@ export default function ModelsAndEndpointsPage() { ? (["retry-settings", "model-group-alias", "access-group-budgets", "price-data"] as const) : []), ], - [canCreate, isAdmin, isViewOnly], + [canCreate, canViewAutoRouters, isAdmin, isViewOnly], ); const allModelsLabel = isAdmin ? "All Models" : "Your Models"; @@ -165,7 +170,9 @@ export default function ModelsAndEndpointsPage() { {isAdmin ? (

Add and manage models for the proxy

) : ( -

Add models for teams you are an admin for.

+

+ View your models and manage routers for teams that allow it. +

)} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx index 12f0b95bf13..1b4251c7ac0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx @@ -12,10 +12,12 @@ vi.mock("../components/AutoRouters/AutoRoutersPanel", () => ({ })); const mockUseAuthorized = vi.fn(); +const mockUseTeams = vi.fn().mockReturnValue({ data: [] }); +const mockUseUISettings = vi.fn(() => ({ data: { values: {} } })); vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => mockUseAuthorized() })); -vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => ({ data: [] }) })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => mockUseTeams() })); vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ - useUISettings: () => ({ data: { values: {} } }), + useUISettings: () => mockUseUISettings(), })); const SESSION = { accessToken: "at", userRole: "Admin", userId: "u1", isViewOnly: false }; @@ -23,6 +25,23 @@ const SESSION = { accessToken: "at", userRole: "Admin", userId: "u1", isViewOnly const lastProps = () => panelProps.mock.calls.at(-1)?.[0] as { createScope: string }; describe("AutoRoutersTabPanel", () => { + it("honors member auto-router opt-in when general model creation is disabled", () => { + mockUseAuthorized.mockReturnValue({ ...SESSION, userRole: "Internal User" }); + mockUseTeams.mockReturnValueOnce({ + data: [ + { + team_id: "team-1", + members_with_roles: [{ user_id: "u1", role: "user" }], + team_member_permissions: ["/auto_router/manage"], + }, + ], + }); + mockUseUISettings.mockReturnValueOnce({ data: { values: { disable_model_add_for_internal_users: true } } }); + render(); + + expect(lastProps().createScope).toBe("team-required"); + }); + it("grants an unscoped create to a real proxy admin", () => { mockUseAuthorized.mockReturnValue(SESSION); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx index 69b442da09b..260b463241b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx @@ -4,14 +4,13 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { internalUserRoles } from "@/utils/roles"; -import { modelCreationScope } from "@/utils/modelPermissions"; +import { autoRouterCreationScope } from "@/utils/modelPermissions"; import { AutoRoutersPanel } from "../components/AutoRouters/AutoRoutersPanel"; /** * Owns the permission decision for the Auto-Routers tab so the panel stays a renderer. - * Creating an auto router is a POST /model/new, the same endpoint Add Model posts to, so it - * takes the same audience rule: a proxy admin, or a team admin who scopes it to a team. + * Auto routers also admit members of teams that enabled their dedicated management grant. * Viewer roles reach the list without write affordances. */ export default function AutoRoutersTabPanel() { @@ -20,7 +19,7 @@ export default function AutoRoutersTabPanel() { const { data: uiSettings } = useUISettings(); const isInternalUser = userRole != null && internalUserRoles.includes(userRole); - const scope = modelCreationScope( + const scope = autoRouterCreationScope( { userRole, userID, isViewOnly }, { teams: teams ?? null, diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index f6e619c5e84..663a608516c 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -15,7 +15,7 @@ import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; import { getMissingTiersError } from "./build_complexity_router_config"; import { getSubmitBlockedReason } from "./add_auto_router_tab"; import { buildModelAvailability } from "@/lib/autorouter_presets"; -import { testAutoRouterRouting } from "../networking"; +import { modelCreateCall, testAutoRouterRouting } from "../networking"; import { ModelGroup } from "@/components/llm_calls/fetch_models"; import { AutoRouterPreset, getRequiredModelsInPreset } from "@/lib/autorouter_presets"; import { BUNDLED_PRESETS, LOADED_PRESETS_QUERY, useAutoRouterPresets } from "../../../tests/mocks/autoRouterPresets"; @@ -104,6 +104,7 @@ const { validateAutoRouterConfig } = vi.hoisted(() => ({ })); vi.mock("../networking", () => ({ + modelCreateCall: vi.fn().mockResolvedValue({}), modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), testAutoRouterRouting: vi.fn(), validateAutoRouterConfig, @@ -111,6 +112,7 @@ vi.mock("../networking", () => ({ vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: mockFetchAvailableModels, + fetchAutoRouterModels: mockFetchAvailableModels, })); vi.mock("@/app/(dashboard)/hooks/models/useModels", async (importOriginal) => { @@ -143,6 +145,7 @@ vi.mock("../common_components/team_dropdown", () => ({ > + @@ -876,6 +826,7 @@ export default function ModelInfoView({ modelData={localModelData || modelData} accessToken={accessToken || ""} userRole={userRole || ""} + isMemberManaged={!canEditModel} /> !open && setIsAutoRouterTestModalOpen(false)}> diff --git a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx index 68ca9d8a2cb..a946f9c2b3d 100644 --- a/ui/litellm-dashboard/src/components/team/permission_definitions.tsx +++ b/ui/litellm-dashboard/src/components/team/permission_definitions.tsx @@ -9,6 +9,7 @@ export interface PermissionInfo { * Map of permission endpoint patterns to their descriptions */ export const PERMISSION_DESCRIPTIONS: Record = { + "/auto_router/manage": "Member can create auto routers for this team and edit their own router configurations", "/key/generate": "Member can generate a virtual key for this team", "/key/service-account/generate": "Member can generate a service account key (not belonging to any user) for this team", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8aa05cf8c7c..5573093d72c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1490,8 +1490,8 @@ export interface paths { * * Runs the same check every write path runs (the router's own pydantic model), so a form can * show the backend's exact verdict while the operator is still editing rather than after a - * rejected save. Gated exactly like the save it rehearses: a proxy admin, or a team admin - * naming their own team. Nothing is created, routed, or billed. + * rejected save. Uses the same team opt-in and model-access checks as configuration + * writes for members. Nothing is created, routed, or billed. */ post: operations["validate_complexity_router_config_auto_router_validate_complexity_router_config_post"]; delete?: never; @@ -28577,7 +28577,7 @@ export interface components { * @description Enum for key management routes * @enum {string} */ - KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/spend/logs" | "/spend/logs/v2"; + KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/spend/logs" | "/spend/logs/v2"; /** * KeyManagementSystem * @enum {string} @@ -39886,6 +39886,11 @@ export interface components { input_cost_per_token?: number | null; /** Internal Router Model */ internal_router_model?: boolean | null; + /** + * Member Auto Router + * @default false + */ + member_auto_router: boolean; /** Output Cost Per Character */ output_cost_per_character?: number | null; /** Output Cost Per Token */ diff --git a/ui/litellm-dashboard/src/utils/modelPermissions.test.ts b/ui/litellm-dashboard/src/utils/modelPermissions.test.ts index afc2ecd210f..900bbbfa01d 100644 --- a/ui/litellm-dashboard/src/utils/modelPermissions.test.ts +++ b/ui/litellm-dashboard/src/utils/modelPermissions.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import { Team } from "@/components/networking"; -import { canCreateModels, canModifyModel, modelCreationScope } from "./modelPermissions"; +import { canCreateModels, canEditAutoRouter, canModifyModel, modelCreationScope } from "./modelPermissions"; const teamWhere = (userId: string, role: string, teamId = "team-1"): Team[] => [{ team_id: teamId, members_with_roles: [{ user_id: userId, user_email: "t@test.com", role }] }] as unknown as Team[]; @@ -111,3 +111,32 @@ describe("canModifyModel", () => { expect(canModifyModel(VIEW_ONLY_ADMIN, teamWhere("u-viewer", "admin"), teamRow)).toBe(false); }); }); + +describe("team member auto routers", () => { + const team = { ...teamWhere("u-member", "user")[0], team_member_permissions: ["/auto_router/manage"] }; + const ownRouter = { + teamId: "team-1", + isDbModel: true, + createdBy: "u-member", + model: "auto_router/complexity_router", + }; + + const revokedTeam: Team = { ...team, team_member_permissions: [] }; + const removedMemberTeam: Team = { ...team, members_with_roles: [] }; + + it.each([ + ["creator", MEMBER, team, ownRouter, true], + ["peer", MEMBER, team, { ...ownRouter, createdBy: "peer" }, false], + ["foreign team", MEMBER, team, { ...ownRouter, teamId: "other-team" }, false], + ["missing creator", MEMBER, team, { ...ownRouter, createdBy: null }, false], + ["config deployment", MEMBER, team, { ...ownRouter, isDbModel: false }, false], + ["ordinary model", MEMBER, team, { ...ownRouter, model: "openai/gpt-5" }, false], + ["revoked permission", MEMBER, revokedTeam, ownRouter, false], + ["removed member", MEMBER, removedMemberTeam, ownRouter, false], + ["blocked team", MEMBER, { ...team, blocked: true }, ownRouter, false], + ["viewer", { ...MEMBER, isViewOnly: true }, team, ownRouter, false], + ] as const)("allows own configuration edits only: %s", (...args) => { + const [, actor, eligibleTeam, origin, expected] = args; + expect(canEditAutoRouter(actor, [eligibleTeam], origin)).toBe(expected); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/modelPermissions.ts b/ui/litellm-dashboard/src/utils/modelPermissions.ts index 843d9041026..7e0a07e8e46 100644 --- a/ui/litellm-dashboard/src/utils/modelPermissions.ts +++ b/ui/litellm-dashboard/src/utils/modelPermissions.ts @@ -11,9 +11,8 @@ import { isProxyAdminRole, isUserTeamAdminForAnyTeam, isUserTeamAdminForSingleTe * * Past that route gate, both questions below are answered by exactly two inputs: the * caller's role, and whether the caller admins the team named in `model_info.team_id`. - * `created_by` is written at creation and never read by an auth check, so it is deliberately - * absent here; gating on it hid controls from team admins the API accepts, and showed - * controls to former team admins the API rejects. + * General model management depends on team administration. The separate auto-router + * member grant below also requires the stored creator for configuration updates. */ export interface ModelActor { userRole: string | null; @@ -95,3 +94,39 @@ export const canModifyModel = ( } return isTeamAdminOf(teams, actor.userID, teamId); }; + +const canMemberCreateAutoRouterForTeam = (actor: ModelActor, team: Team): boolean => { + if (actor.isViewOnly || !actor.userID) return false; + const membership = team.members_with_roles.find((member) => member.user_id === actor.userID); + return ( + membership?.role === "user" && + !team.blocked && + team.team_member_permissions?.includes("/auto_router/manage") === true + ); +}; + +export const canCreateAutoRouterForTeam = (actor: ModelActor, team: Team): boolean => { + if (actor.isViewOnly || !actor.userID) return false; + return ( + canModifyModel(actor, [team], { teamId: team.team_id, isDbModel: true }) || + canMemberCreateAutoRouterForTeam(actor, team) + ); +}; + +export const autoRouterCreationScope = (actor: ModelActor, limits: ModelCreationLimits): ModelWriteScope => { + const scope = modelCreationScope(actor, limits); + if (scope !== "forbidden") return scope; + return limits.teams?.some((team) => canMemberCreateAutoRouterForTeam(actor, team)) ? "team-required" : "forbidden"; +}; + +export const canEditAutoRouter = ( + actor: ModelActor, + teams: Team[] | null, + origin: ModelRowOrigin & { createdBy: string | null | undefined; model: string | null | undefined }, +): boolean => { + if (canModifyModel(actor, teams, origin)) return true; + if (!origin.isDbModel || !actor.userID) return false; + if (actor.userID !== origin.createdBy || origin.model !== "auto_router/complexity_router") return false; + const team = teams?.find((candidate) => candidate.team_id === origin.teamId); + return team != null && canCreateAutoRouterForTeam(actor, team); +};