diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl new file mode 100644 index 0000000000..c98d9cfcfa Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz new file mode 100644 index 0000000000..c8c3340462 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql new file mode 100644 index 0000000000..67e75e84c4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql @@ -0,0 +1,33 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- AlterTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- CreateTable +CREATE TABLE "LiteLLM_AccessGroupTable" ( + "access_group_id" TEXT NOT NULL, + "access_group_name" TEXT NOT NULL, + "description" TEXT, + "access_model_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], + "access_mcp_server_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], + "access_agent_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], + "assigned_team_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], + "assigned_key_ids" TEXT[] DEFAULT ARRAY[]::TEXT[], + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_AccessGroupTable_pkey" PRIMARY KEY ("access_group_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_AccessGroupTable_access_group_name_key" ON "LiteLLM_AccessGroupTable"("access_group_name"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 558dfcc951..2a0faee280 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -128,6 +128,7 @@ model LiteLLM_TeamTable { model_max_budget Json @default("{}") router_settings Json? @default("{}") team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) policies String[] @default([]) model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team @@ -161,6 +162,7 @@ model LiteLLM_DeletedTeamTable { model_max_budget Json @default("{}") router_settings Json? @default("{}") team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) policies String[] @default([]) model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) @@ -293,6 +295,7 @@ model LiteLLM_VerificationToken { allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) policies String[] @default([]) + access_group_ids String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_id String? @@ -348,6 +351,7 @@ model LiteLLM_DeletedVerificationToken { allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) policies String[] @default([]) + access_group_ids String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") router_settings Json? @default("{}") @@ -920,3 +924,23 @@ model LiteLLM_PolicyAttachmentTable { updated_at DateTime @default(now()) @updatedAt updated_by String? } + +//Unified Access Groups table for storing unified access groups +model LiteLLM_AccessGroupTable { + access_group_id String @id @default(uuid()) + access_group_name String @unique + description String? + + // Resource memberships - explicit arrays per type + access_model_ids String[] @default([]) + access_mcp_server_ids String[] @default([]) + access_agent_ids String[] @default([]) + + assigned_team_ids String[] @default([]) + assigned_key_ids String[] @default([]) + + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} \ No newline at end of file diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index ceeb8c8bb4..eda49bfb9f 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm-proxy-extras" -version = "0.4.35" +version = "0.4.36" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." authors = ["BerriAI"] readme = "README.md" @@ -22,7 +22,7 @@ requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "0.4.35" +version = "0.4.36" version_files = [ "pyproject.toml:version", "../requirements.txt:litellm-proxy-extras==", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 87ff4a66e0..0c85e9ba4f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -893,6 +893,7 @@ class KeyRequestBase(GenerateRequestBase): Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm router_settings: Optional[UpdateRouterConfig] = None + access_group_ids: Optional[List[str]] = None class LiteLLMKeyType(str, enum.Enum): @@ -1502,6 +1503,7 @@ class TeamBase(LiteLLMPydanticObjectBase): models: list = [] blocked: bool = False router_settings: Optional[dict] = None + access_group_ids: Optional[List[str]] = None class NewTeamRequest(TeamBase): @@ -1589,6 +1591,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): model_tpm_limit: Optional[Dict[str, int]] = None allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None router_settings: Optional[dict] = None + access_group_ids: Optional[List[str]] = None class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -2177,6 +2180,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): updated_by: Optional[str] = None object_permission_id: Optional[str] = None object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + access_group_ids: Optional[List[str]] = None rotation_count: Optional[int] = 0 # Number of times key has been rotated auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py new file mode 100644 index 0000000000..33b81e8565 --- /dev/null +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -0,0 +1,264 @@ +from typing import List + +from fastapi import APIRouter, Depends, HTTPException, status + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.utils import get_prisma_client_or_throw +from litellm.types.access_group import ( + AccessGroupCreateRequest, + AccessGroupResponse, + AccessGroupUpdateRequest, +) + +router = APIRouter( + tags=["access group management"], +) + + +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + +def _record_to_response(record) -> AccessGroupResponse: + return AccessGroupResponse( + access_group_id=record.access_group_id, + access_group_name=record.access_group_name, + description=record.description, + access_model_ids=record.access_model_ids, + access_mcp_server_ids=record.access_mcp_server_ids, + access_agent_ids=record.access_agent_ids, + assigned_team_ids=record.assigned_team_ids, + assigned_key_ids=record.assigned_key_ids, + created_at=record.created_at, + created_by=record.created_by, + updated_at=record.updated_at, + updated_by=record.updated_by, + ) + + +@router.post( + "/v1/access_group", + response_model=AccessGroupResponse, + status_code=status.HTTP_201_CREATED, +) +async def create_access_group( + data: AccessGroupCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> AccessGroupResponse: + _require_proxy_admin(user_api_key_dict) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + + existing = await prisma_client.db.litellm_accessgrouptable.find_unique( + where={"access_group_name": data.access_group_name} + ) + if existing is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Access group '{data.access_group_name}' already exists", + ) + + try: + record = await prisma_client.db.litellm_accessgrouptable.create( + data={ + "access_group_name": data.access_group_name, + "description": data.description, + "access_model_ids": data.access_model_ids or [], + "access_mcp_server_ids": data.access_mcp_server_ids or [], + "access_agent_ids": data.access_agent_ids or [], + "assigned_team_ids": data.assigned_team_ids or [], + "assigned_key_ids": data.assigned_key_ids or [], + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + except Exception as e: + # Race condition: another request created the same name between find_unique and create. + # Prisma raises UniqueViolationError (P2002) or similar for unique constraint. + if "unique constraint" in str(e).lower() or "P2002" in str(e): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Access group '{data.access_group_name}' already exists", + ) + raise + return _record_to_response(record) + + +@router.get( + "/v1/access_group", + response_model=List[AccessGroupResponse], +) +async def list_access_groups( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> List[AccessGroupResponse]: + _require_proxy_admin(user_api_key_dict) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + + records = await prisma_client.db.litellm_accessgrouptable.find_many( + order={"created_at": "desc"} + ) + return [_record_to_response(r) for r in records] + + +@router.get( + "/v1/access_group/{access_group_id}", + response_model=AccessGroupResponse, +) +async def get_access_group( + access_group_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> AccessGroupResponse: + _require_proxy_admin(user_api_key_dict) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + + record = await prisma_client.db.litellm_accessgrouptable.find_unique( + where={"access_group_id": access_group_id} + ) + if record is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) + return _record_to_response(record) + + +@router.put( + "/v1/access_group/{access_group_id}", + response_model=AccessGroupResponse, +) +async def update_access_group( + access_group_id: str, + data: AccessGroupUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> AccessGroupResponse: + _require_proxy_admin(user_api_key_dict) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + + existing = await prisma_client.db.litellm_accessgrouptable.find_unique( + where={"access_group_id": access_group_id} + ) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) + + update_data: dict = {"updated_by": user_api_key_dict.user_id} + for field, value in data.model_dump(exclude_unset=True).items(): + update_data[field] = value + + record = await prisma_client.db.litellm_accessgrouptable.update( + where={"access_group_id": access_group_id}, + data=update_data, + ) + return _record_to_response(record) + + +@router.delete( + "/v1/access_group/{access_group_id}", + status_code=status.HTTP_204_NO_CONTENT, +) +async def delete_access_group( + access_group_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> None: + _require_proxy_admin(user_api_key_dict) + prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) + + try: + async with prisma_client.db.tx() as tx: + existing = await tx.litellm_accessgrouptable.find_unique( + where={"access_group_id": access_group_id} + ) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) + + # Remove access_group_id from teams and keys that reference it + teams_with_group = await tx.litellm_teamtable.find_many( + where={"access_group_ids": {"hasSome": [access_group_id]}} + ) + for team in teams_with_group: + updated_ids = [tid for tid in (team.access_group_ids or []) if tid != access_group_id] + await tx.litellm_teamtable.update( + where={"team_id": team.team_id}, + data={"access_group_ids": updated_ids}, + ) + + keys_with_group = await tx.litellm_verificationtoken.find_many( + where={"access_group_ids": {"hasSome": [access_group_id]}} + ) + for key in keys_with_group: + updated_ids = [kid for kid in (key.access_group_ids or []) if kid != access_group_id] + await tx.litellm_verificationtoken.update( + where={"token": key.token}, + data={"access_group_ids": updated_ids}, + ) + + await tx.litellm_accessgrouptable.delete( + where={"access_group_id": access_group_id} + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "delete_access_group failed: access_group_id=%s error=%s", + access_group_id, + e, + ) + if PrismaDBExceptionHandler.is_database_connection_error(e): + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + if "P2025" in str(e) or ("record" in str(e).lower() and "not found" in str(e).lower()): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to delete access group. Please try again.", + ) + + +# Alias routes for /v1/unified_access_group +router.add_api_route( + "/v1/unified_access_group", + create_access_group, + methods=["POST"], + response_model=AccessGroupResponse, + status_code=status.HTTP_201_CREATED, +) +router.add_api_route( + "/v1/unified_access_group", + list_access_groups, + methods=["GET"], + response_model=List[AccessGroupResponse], +) +router.add_api_route( + "/v1/unified_access_group/{access_group_id}", + get_access_group, + methods=["GET"], + response_model=AccessGroupResponse, +) +router.add_api_route( + "/v1/unified_access_group/{access_group_id}", + update_access_group, + methods=["PUT"], + response_model=AccessGroupResponse, +) +router.add_api_route( + "/v1/unified_access_group/{access_group_id}", + delete_access_group, + methods=["DELETE"], + status_code=status.HTTP_204_NO_CONTENT, +) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7acfa70c5a..fd86c5e1d2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -393,6 +393,9 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( from litellm.proxy.management_endpoints.team_callback_endpoints import ( router as team_callback_router, ) +from litellm.proxy.management_endpoints.access_group_endpoints import ( + router as access_group_router, +) from litellm.proxy.management_endpoints.team_endpoints import router as team_router from litellm.proxy.management_endpoints.team_endpoints import ( update_team, @@ -11998,6 +12001,7 @@ app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) app.include_router(agent_endpoints_router) app.include_router(a2a_router) +app.include_router(access_group_router) ######################################################## # MCP Server ######################################################## diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 37ed018266..b3caa46147 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -128,6 +128,7 @@ model LiteLLM_TeamTable { model_max_budget Json @default("{}") router_settings Json? @default("{}") team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) policies String[] @default([]) model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) @@ -160,6 +161,7 @@ model LiteLLM_DeletedTeamTable { model_max_budget Json @default("{}") router_settings Json? @default("{}") team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) policies String[] @default([]) model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases @@ -291,6 +293,7 @@ model LiteLLM_VerificationToken { allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) policies String[] @default([]) + access_group_ids String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_id String? @@ -346,6 +349,7 @@ model LiteLLM_DeletedVerificationToken { allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) policies String[] @default([]) + access_group_ids String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") router_settings Json? @default("{}") @@ -917,3 +921,23 @@ model LiteLLM_PolicyAttachmentTable { updated_at DateTime @default(now()) @updatedAt updated_by String? } + +//Unified Access Groups table for storing unified access groups +model LiteLLM_AccessGroupTable { + access_group_id String @id @default(uuid()) + access_group_name String @unique + description String? + + // Resource memberships - explicit arrays per type + access_model_ids String[] @default([]) + access_mcp_server_ids String[] @default([]) + access_agent_ids String[] @default([]) + + assigned_team_ids String[] @default([]) + assigned_key_ids String[] @default([]) + + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} \ No newline at end of file diff --git a/litellm/types/access_group.py b/litellm/types/access_group.py new file mode 100644 index 0000000000..3a6b75768e --- /dev/null +++ b/litellm/types/access_group.py @@ -0,0 +1,38 @@ +from datetime import datetime +from typing import List, Optional + +from pydantic import BaseModel + + +class AccessGroupCreateRequest(BaseModel): + access_group_name: str + description: Optional[str] = None + access_model_ids: Optional[List[str]] = None + access_mcp_server_ids: Optional[List[str]] = None + access_agent_ids: Optional[List[str]] = None + assigned_team_ids: Optional[List[str]] = None + assigned_key_ids: Optional[List[str]] = None + + +class AccessGroupUpdateRequest(BaseModel): + description: Optional[str] = None + access_model_ids: Optional[List[str]] = None + access_mcp_server_ids: Optional[List[str]] = None + access_agent_ids: Optional[List[str]] = None + assigned_team_ids: Optional[List[str]] = None + assigned_key_ids: Optional[List[str]] = None + + +class AccessGroupResponse(BaseModel): + access_group_id: str + access_group_name: str + description: Optional[str] = None + access_model_ids: List[str] + access_mcp_server_ids: List[str] + access_agent_ids: List[str] + assigned_team_ids: List[str] + assigned_key_ids: List[str] + created_at: datetime + created_by: Optional[str] = None + updated_at: datetime + updated_by: Optional[str] = None diff --git a/pyproject.toml b/pyproject.toml index 8ba7119ec5..87322be61d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.35", optional = true} +litellm-proxy-extras = {version = "0.4.36", optional = true} rich = {version = "13.7.1", optional = true} litellm-enterprise = {version = "0.1.31", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index 4db8605490..37c97be167 100644 --- a/requirements.txt +++ b/requirements.txt @@ -55,7 +55,7 @@ sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests cryptography==44.0.1 tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.35 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.36 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/schema.prisma b/schema.prisma index 4329f939a7..2a11d0028f 100644 --- a/schema.prisma +++ b/schema.prisma @@ -128,6 +128,7 @@ model LiteLLM_TeamTable { model_max_budget Json @default("{}") router_settings Json? @default("{}") team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) policies String[] @default([]) model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team @@ -161,6 +162,7 @@ model LiteLLM_DeletedTeamTable { model_max_budget Json @default("{}") router_settings Json? @default("{}") team_member_permissions String[] @default([]) + access_group_ids String[] @default([]) policies String[] @default([]) model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) @@ -293,6 +295,7 @@ model LiteLLM_VerificationToken { allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) policies String[] @default([]) + access_group_ids String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") budget_id String? @@ -348,6 +351,7 @@ model LiteLLM_DeletedVerificationToken { allowed_cache_controls String[] @default([]) allowed_routes String[] @default([]) policies String[] @default([]) + access_group_ids String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") router_settings Json? @default("{}") @@ -363,7 +367,6 @@ model LiteLLM_DeletedVerificationToken { rotation_interval String? last_rotation_at DateTime? key_rotation_at DateTime? - // Deletion metadata deleted_at DateTime @default(now()) @map("deleted_at") deleted_by String? @map("deleted_by") // User who deleted the key @@ -919,3 +922,23 @@ model LiteLLM_PolicyAttachmentTable { updated_at DateTime @default(now()) @updatedAt updated_by String? } + +//Unified Access Groups table for storing unified access groups +model LiteLLM_AccessGroupTable { + access_group_id String @id @default(uuid()) + access_group_name String @unique + description String? + + // Resource memberships - explicit arrays per type + access_model_ids String[] @default([]) + access_mcp_server_ids String[] @default([]) + access_agent_ids String[] @default([]) + + assigned_team_ids String[] @default([]) + assigned_key_ids String[] @default([]) + + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} \ No newline at end of file diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py new file mode 100644 index 0000000000..5f204918d0 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -0,0 +1,583 @@ +""" +Tests for access group management endpoints. +""" + +import os +import sys +import types +from contextlib import asynccontextmanager +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient +from prisma.errors import PrismaError + +import litellm.proxy.proxy_server as ps +from litellm.proxy.proxy_server import app +from litellm.proxy._types import ( + CommonProxyErrors, + LitellmUserRoles, + UserAPIKeyAuth, +) + +sys.path.insert(0, os.path.abspath("../../../")) + + +def _make_access_group_record( + access_group_id: str = "ag-123", + access_group_name: str = "test-group", + description: str | None = "Test description", + access_model_ids: list | None = None, + access_mcp_server_ids: list | None = None, + access_agent_ids: list | None = None, + assigned_team_ids: list | None = None, + assigned_key_ids: list | None = None, + created_by: str | None = "admin-user", + updated_by: str | None = "admin-user", + created_at: datetime | None = None, +): + record = MagicMock() + record.access_group_id = access_group_id + record.access_group_name = access_group_name + record.description = description + record.access_model_ids = access_model_ids or [] + record.access_mcp_server_ids = access_mcp_server_ids or [] + record.access_agent_ids = access_agent_ids or [] + record.assigned_team_ids = assigned_team_ids or [] + record.assigned_key_ids = assigned_key_ids or [] + record.created_at = created_at or datetime.now() + record.created_by = created_by + record.updated_at = datetime.now() + record.updated_by = updated_by + return record + + +@pytest.fixture +def client_and_mocks(monkeypatch): + """Setup mock prisma and admin auth for access group endpoints.""" + mock_access_group_table = MagicMock() + mock_prisma = MagicMock() + + def _create_side_effect(*, data): + return _make_access_group_record( + access_group_id="ag-new", + access_group_name=data.get("access_group_name", "new"), + description=data.get("description"), + access_model_ids=data.get("access_model_ids", []), + access_mcp_server_ids=data.get("access_mcp_server_ids", []), + access_agent_ids=data.get("access_agent_ids", []), + assigned_team_ids=data.get("assigned_team_ids", []), + assigned_key_ids=data.get("assigned_key_ids", []), + created_by=data.get("created_by"), + updated_by=data.get("updated_by"), + ) + + mock_access_group_table.create = AsyncMock(side_effect=_create_side_effect) + mock_access_group_table.find_unique = AsyncMock(return_value=None) + mock_access_group_table.find_many = AsyncMock(return_value=[]) + mock_access_group_table.update = AsyncMock(side_effect=lambda *, where, data: _make_access_group_record( + access_group_id=where.get("access_group_id", "ag-123"), + access_group_name=data.get("access_group_name", "updated"), + description=data.get("description"), + access_model_ids=data.get("access_model_ids", []), + access_mcp_server_ids=data.get("access_mcp_server_ids", []), + access_agent_ids=data.get("access_agent_ids", []), + assigned_team_ids=data.get("assigned_team_ids", []), + assigned_key_ids=data.get("assigned_key_ids", []), + updated_by=data.get("updated_by"), + )) + mock_access_group_table.delete = AsyncMock(return_value=None) + + mock_team_table = MagicMock() + mock_team_table.find_many = AsyncMock(return_value=[]) + mock_team_table.update = AsyncMock(return_value=None) + + mock_key_table = MagicMock() + mock_key_table.find_many = AsyncMock(return_value=[]) + mock_key_table.update = AsyncMock(return_value=None) + + @asynccontextmanager + async def mock_tx(): + tx = types.SimpleNamespace( + litellm_accessgrouptable=mock_access_group_table, + litellm_teamtable=mock_team_table, + litellm_verificationtoken=mock_key_table, + ) + yield tx + + mock_db = types.SimpleNamespace( + litellm_accessgrouptable=mock_access_group_table, + litellm_teamtable=mock_team_table, + litellm_verificationtoken=mock_key_table, + tx=mock_tx, + ) + mock_prisma.db = mock_db + + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + admin_user = UserAPIKeyAuth( + user_id="admin_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: admin_user + + client = TestClient(app) + + yield client, mock_prisma, mock_access_group_table + + app.dependency_overrides.clear() + monkeypatch.setattr(ps, "prisma_client", ps.prisma_client) + + +# Paths for primary and alias endpoints (alias: /v1/unified_access_group) +ACCESS_GROUP_PATHS = ["/v1/access_group", "/v1/unified_access_group"] + + +# --------------------------------------------------------------------------- +# CREATE +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +@pytest.mark.parametrize( + "payload", + [ + {"access_group_name": "group-a"}, + { + "access_group_name": "group-b", + "description": "Group B description", + "access_model_ids": ["model-1"], + "access_mcp_server_ids": ["mcp-1"], + "assigned_team_ids": ["team-1"], + }, + ], +) +def test_create_access_group_success(client_and_mocks, base_path, payload): + """Create access group with various payloads returns 201.""" + client, _, mock_table = client_and_mocks + + resp = client.post(base_path, json=payload) + assert resp.status_code == 201 + body = resp.json() + assert body["access_group_name"] == payload["access_group_name"] + assert body.get("access_group_id") is not None + mock_table.create.assert_awaited_once() + + +def test_create_access_group_duplicate_name_conflict(client_and_mocks): + """Create with duplicate name returns 409.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_name="existing-group") + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.post("/v1/access_group", json={"access_group_name": "existing-group"}) + assert resp.status_code == 409 + assert "already exists" in resp.json()["detail"] + + +@pytest.mark.parametrize( + "error_message", + [ + "Unique constraint failed on the fields: (`access_group_name`)", + "P2002: Unique constraint failed", + "unique constraint violation", + ], +) +def test_create_access_group_race_condition_returns_409(client_and_mocks, error_message): + """Create race condition: Prisma unique constraint surfaces as 409, not 500.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + mock_table.create = AsyncMock(side_effect=Exception(error_message)) + + resp = client.post("/v1/access_group", json={"access_group_name": "race-group"}) + assert resp.status_code == 409 + assert "already exists" in resp.json()["detail"] + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_create_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot create access groups.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.post("/v1/access_group", json={"access_group_name": "forbidden"}) + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +def test_create_access_group_validation_missing_name(client_and_mocks): + """Create with missing access_group_name returns 422.""" + client, _, _ = client_and_mocks + + resp = client.post("/v1/access_group", json={}) + assert resp.status_code == 422 + + +def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks): + """Create with non-unique-constraint Prisma error returns 500.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + mock_table.create = AsyncMock(side_effect=Exception("Some other database error")) + + # Use raise_server_exceptions=False so unhandled exceptions become 500 responses + test_client = TestClient(app, raise_server_exceptions=False) + resp = test_client.post("/v1/access_group", json={"access_group_name": "test-group"}) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# LIST +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_list_access_groups_success_empty(client_and_mocks, base_path): + """List access groups returns empty list when none exist.""" + client, _, mock_table = client_and_mocks + + resp = client.get(base_path) + assert resp.status_code == 200 + assert resp.json() == [] + mock_table.find_many.assert_awaited_once() + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_list_access_groups_success_with_items(client_and_mocks, base_path): + """List access groups returns items when they exist.""" + client, _, mock_table = client_and_mocks + + records = [ + _make_access_group_record(access_group_id="ag-1", access_group_name="group-1"), + _make_access_group_record(access_group_id="ag-2", access_group_name="group-2"), + ] + mock_table.find_many = AsyncMock(return_value=records) + + resp = client.get(base_path) + assert resp.status_code == 200 + body = resp.json() + assert len(body) == 2 + assert body[0]["access_group_name"] == "group-1" + assert body[1]["access_group_name"] == "group-2" + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_list_access_groups_ordered_by_created_at_desc(client_and_mocks, base_path): + """List access groups calls find_many with created_at desc order.""" + client, _, mock_table = client_and_mocks + + older = datetime(2025, 1, 1, 12, 0, 0) + newer = datetime(2025, 1, 2, 12, 0, 0) + records = [ + _make_access_group_record( + access_group_id="ag-newer", + access_group_name="newer-group", + created_at=newer, + ), + _make_access_group_record( + access_group_id="ag-older", + access_group_name="older-group", + created_at=older, + ), + ] + mock_table.find_many = AsyncMock(return_value=records) + + resp = client.get(base_path) + assert resp.status_code == 200 + body = resp.json() + assert len(body) == 2 + # Mock returns newest first (simulating Prisma order desc) + assert body[0]["access_group_name"] == "newer-group" + assert body[1]["access_group_name"] == "older-group" + mock_table.find_many.assert_awaited_once_with(order={"created_at": "desc"}) + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_list_access_groups_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot list access groups.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.get("/v1/access_group") + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# GET +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +@pytest.mark.parametrize("access_group_id", ["ag-123", "ag-other-id"]) +def test_get_access_group_success(client_and_mocks, base_path, access_group_id): + """Get access group by id returns record when found.""" + client, _, mock_table = client_and_mocks + + record = _make_access_group_record(access_group_id=access_group_id) + mock_table.find_unique = AsyncMock(return_value=record) + + resp = client.get(f"{base_path}/{access_group_id}") + assert resp.status_code == 200 + assert resp.json()["access_group_id"] == access_group_id + + +def test_get_access_group_not_found(client_and_mocks): + """Get access group returns 404 when not found.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.get("/v1/access_group/nonexistent-id") + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_get_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot get access group.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.get("/v1/access_group/ag-123") + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# UPDATE +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +@pytest.mark.parametrize( + "update_payload", + [ + {"description": "Updated description"}, + {"access_model_ids": ["model-1", "model-2"]}, + {"assigned_team_ids": [], "assigned_key_ids": ["key-1"]}, + ], +) +def test_update_access_group_success(client_and_mocks, base_path, update_payload): + """Update access group with various payloads returns 200.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-update") + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.put(f"{base_path}/ag-update", json=update_payload) + assert resp.status_code == 200 + mock_table.update.assert_awaited_once() + + +def test_update_access_group_not_found(client_and_mocks): + """Update access group returns 404 when not found.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.put( + "/v1/access_group/nonexistent-id", + json={"description": "Updated"}, + ) + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + mock_table.update.assert_not_awaited() + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_update_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot update access groups.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.put("/v1/access_group/ag-123", json={"description": "Updated"}) + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +def test_update_access_group_empty_body(client_and_mocks): + """Update with empty body succeeds; only updated_by is set.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="unchanged") + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.put("/v1/access_group/ag-update", json={}) + assert resp.status_code == 200 + mock_table.update.assert_awaited_once() + call_kwargs = mock_table.update.call_args.kwargs + assert call_kwargs["where"] == {"access_group_id": "ag-update"} + assert "updated_by" in call_kwargs["data"] + assert call_kwargs["data"]["updated_by"] == "admin_user" + + +# --------------------------------------------------------------------------- +# DELETE +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +@pytest.mark.parametrize("access_group_id", ["ag-123", "ag-delete-me"]) +def test_delete_access_group_success(client_and_mocks, base_path, access_group_id): + """Delete access group returns 204 when found.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id=access_group_id) + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.delete(f"{base_path}/{access_group_id}") + assert resp.status_code == 204 + mock_table.delete.assert_awaited_once() + + +def test_delete_access_group_not_found(client_and_mocks): + """Delete access group returns 404 when not found.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.delete("/v1/access_group/nonexistent-id") + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + mock_table.delete.assert_not_awaited() + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot delete access groups.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.delete("/v1/access_group/ag-123") + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): + """Delete removes access_group_id from teams and keys before deleting the group.""" + client, mock_prisma, mock_access_group_table = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + team_with_group = MagicMock() + team_with_group.team_id = "team-1" + team_with_group.access_group_ids = ["ag-to-delete", "ag-other"] + mock_team_table.find_many = AsyncMock(return_value=[team_with_group]) + + key_with_group = MagicMock() + key_with_group.token = "key-token-1" + key_with_group.access_group_ids = ["ag-to-delete"] + mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_team_table.update.assert_awaited_once_with( + where={"team_id": "team-1"}, + data={"access_group_ids": ["ag-other"]}, + ) + mock_key_table.update.assert_awaited_once_with( + where={"token": "key-token-1"}, + data={"access_group_ids": []}, + ) + mock_access_group_table.delete.assert_awaited_once_with( + where={"access_group_id": "ag-to-delete"} + ) + + +def test_delete_access_group_503_on_db_connection_error(client_and_mocks): + """Delete returns 503 when DB connection error occurs during transaction.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_table.find_unique = AsyncMock(return_value=existing) + mock_table.delete = AsyncMock(side_effect=PrismaError()) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 503 + assert resp.json()["detail"] == CommonProxyErrors.db_not_connected_error.value + + +def test_delete_access_group_404_on_p2025_or_record_not_found(client_and_mocks): + """Delete returns 404 when Prisma raises P2025 or record-not-found error.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_table.find_unique = AsyncMock(return_value=existing) + mock_table.delete = AsyncMock(side_effect=Exception("P2025: Record to delete does not exist")) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + + +def test_delete_access_group_500_on_generic_exception(client_and_mocks): + """Delete returns 500 when generic exception occurs during transaction.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_table.find_unique = AsyncMock(return_value=existing) + mock_table.delete = AsyncMock(side_effect=RuntimeError("Unexpected error")) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 500 + assert "Failed to delete access group" in resp.json()["detail"] + + +# --------------------------------------------------------------------------- +# DB NOT CONNECTED +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "method,url,factory", + [ + ("post", "/v1/access_group", lambda: {"json": {"access_group_name": "test"}}), + ("get", "/v1/access_group", lambda: {}), + ("get", "/v1/access_group/ag-123", lambda: {}), + ("put", "/v1/access_group/ag-123", lambda: {"json": {"description": "x"}}), + ("delete", "/v1/access_group/ag-123", lambda: {}), + # Alias: /v1/unified_access_group + ("post", "/v1/unified_access_group", lambda: {"json": {"access_group_name": "test"}}), + ("get", "/v1/unified_access_group", lambda: {}), + ("get", "/v1/unified_access_group/ag-123", lambda: {}), + ("put", "/v1/unified_access_group/ag-123", lambda: {"json": {"description": "x"}}), + ("delete", "/v1/unified_access_group/ag-123", lambda: {}), + ], +) +def test_access_group_endpoints_db_not_connected(client_and_mocks, monkeypatch, method, url, factory): + """All endpoints return 500 when DB is not connected.""" + client, _, _ = client_and_mocks + + monkeypatch.setattr(ps, "prisma_client", None) + + resp = getattr(client, method)(url, **factory()) + assert resp.status_code == 500 + assert resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value