diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 02574ca505db..f3a028f58055 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -18,6 +18,7 @@ "/team/", "/v2/team/", "/organization/", + "/v2/organization/", "/customer/", "/end_user/", "/sso/", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 20dc14aa5a38..b73841c47938 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2778,6 +2778,30 @@ def set_model_info(cls, values): return values +class OrganizationUpdateRequestV2(LiteLLMPydanticObjectBase): + """ + Typed PATCH body for ``/v2/organization/{organization_id}`` (RFC 7396 merge-patch). + + Presence is read from ``model_fields_set``, so a sent field is written and an omitted one is + left untouched. ``extra="forbid"`` makes an unknown key a 422 rather than a silent no-op, since + the contract hinges on which keys are present. See the endpoint for the per-field clear tokens. + """ + + model_config = ConfigDict(extra="forbid") + + organization_alias: str | None = None + models: list[str] | None = None + metadata: dict | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + max_budget: float | None = None + soft_budget: float | None = None + max_parallel_requests: int | None = None + model_max_budget: dict | None = None + budget_duration: str | None = None + object_permission: LiteLLM_ObjectPermissionBase | None = None + + from litellm.models.organization import ( # noqa: E402 LiteLLM_OrganizationTable as LiteLLM_OrganizationTable, ) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 138a55d92277..5a289d22f993 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -13,16 +13,18 @@ #### ORGANIZATION MANAGEMENT #### -from typing import Any, Dict, List, Optional, Tuple +from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, status +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import can_user_call_model, get_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, @@ -34,6 +36,7 @@ ) from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, + prepare_object_permission_upsert, ) from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, @@ -101,6 +104,30 @@ async def _verify_org_access( ) +_STR_OBJECT_DICT_ADAPTER = TypeAdapter(dict[str, object]) +_BUDGET_SETTABLE_FIELDS = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"} +_ORG_COLUMN_FIELDS = frozenset({"organization_alias", "models"}) + + +def build_budget_write_data(budget_updates: Mapping[str, object], updated_by: str) -> Mapping[str, object]: + """ + Budget-row columns to write. ``budget_reset_at`` tracks any sent ``budget_duration``: + recomputed for a new duration, cleared alongside a ``None`` duration so no stale reset + timestamp survives. Other sent fields (including a ``None`` clear) are written as-is. + """ + budget_duration = budget_updates.get("budget_duration") + recomputed_reset_at: Mapping[str, object] = ( + { + "budget_reset_at": ( + get_budget_reset_time(budget_duration=budget_duration) if isinstance(budget_duration, str) else None + ) + } + if "budget_duration" in budget_updates + else {} + ) + return {**budget_updates, **recomputed_reset_at, "updated_by": updated_by} + + def handle_nested_budget_structure_in_organization_update_request( raw_data: dict, ) -> dict: @@ -556,6 +583,154 @@ async def handle_update_object_permission( return data_json +@router.patch( + "/v2/organization/{organization_id}", + tags=["organization management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_OrganizationTableWithMembers, + include_in_schema=False, +) +async def update_organization_v2( + organization_id: str, + data: OrganizationUpdateRequestV2, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Partial update of an organization (RESTful PATCH, RFC 7396 merge-patch semantics). + + A sent field is written and an omitted one is left untouched (presence is read from + ``model_fields_set``). Clear tokens are per field: budget limits and ``metadata`` clear with + ``null``, ``models`` with ``[]``, and ``object_permission`` with ``null`` (it merges when sent, + so an empty ``{}`` is rejected). ``organization_alias`` is required and cannot be cleared. + Validation failures return 422; the object-permission upsert, budget-row write, and + org-row write are one transaction. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_id is None: + raise HTTPException( + status_code=400, + detail={ + "error": "Cannot associate a user_id to this action. Check `/key/info` to validate if 'user_id' is set." + }, + ) + + if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0): + raise HTTPException( + status_code=422, + detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"}, + ) + if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0): + raise HTTPException( + status_code=422, + detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, + ) + if data.model_max_budget: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_model_max_budget, + ) + + try: + validate_model_max_budget(data.model_max_budget) + except ValueError as e: + raise HTTPException(status_code=422, detail={"error": str(e)}) + + if "organization_alias" in data.model_fields_set and data.organization_alias is None: + raise HTTPException( + status_code=422, + detail={"error": "organization_alias cannot be cleared; it is required"}, + ) + if "models" in data.model_fields_set and data.models is None: + raise HTTPException( + status_code=422, + detail={"error": "models cannot be set to null; send [] to clear it"}, + ) + if data.object_permission is not None and not data.object_permission.model_dump(exclude_none=True): + raise HTTPException( + status_code=422, + detail={ + "error": "object_permission cannot be an empty object; send null to clear it, or a non-empty object to set grants" + }, + ) + + await _verify_org_access( + organization_id=organization_id, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + + existing_organization_row = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": organization_id}, + ) + if existing_organization_row is None: + raise HTTPException( + status_code=404, + detail={"error": f"Organization not found for organization_id={organization_id}"}, + ) + + field_values = _STR_OBJECT_DICT_ADAPTER.validate_python(data.model_dump()) + present_fields = data.model_fields_set + budget_updates = {field: field_values[field] for field in present_fields if field in _BUDGET_SETTABLE_FIELDS} + org_column_updates: Mapping[str, object] = { + **{field: field_values[field] for field in present_fields if field in _ORG_COLUMN_FIELDS}, + **({"metadata": data.metadata or {}} if "metadata" in present_fields else {}), + } + + object_permission_cleared = "object_permission" in present_fields and data.object_permission is None + object_permission_upsert = ( + await prepare_object_permission_upsert( + new_object_permission=data.object_permission.model_dump(exclude_none=True), + existing_object_permission_id=existing_organization_row.object_permission_id, + prisma_client=prisma_client, + ) + if data.object_permission is not None + else None + ) + object_permission_write: Mapping[str, object] = ( + {"object_permission_id": object_permission_upsert.object_permission_id} + if object_permission_upsert is not None + else ({"object_permission_id": None} if object_permission_cleared else {}) + ) + + organization_write_data = prisma_client.jsonify_object( + { + **org_column_updates, + **object_permission_write, + "updated_by": user_api_key_dict.user_id, + } + ) + + async with prisma_client.db.tx() as tx: + if object_permission_upsert is not None: + await tx.litellm_objectpermissiontable.upsert( + where={"object_permission_id": object_permission_upsert.object_permission_id}, + data={ + "create": object_permission_upsert.record, + "update": object_permission_upsert.record, + }, + ) + if budget_updates: + await tx.litellm_budgettable.update( + where={"budget_id": existing_organization_row.budget_id}, + data=prisma_client.jsonify_object( + dict(build_budget_write_data(budget_updates, user_api_key_dict.user_id)) + ), + ) + response = await tx.litellm_organizationtable.update( + where={"organization_id": organization_id}, + data=organization_write_data, + include={"members": True, "teams": True, "litellm_budget_table": True}, + ) + + return response + + @router.delete( "/organization/delete", tags=["organization management"], diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 6bbd41b93ed2..015ab1d6df6d 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -4,7 +4,8 @@ """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Set, Union from fastapi import HTTPException, status @@ -64,6 +65,57 @@ async def attach_object_permission_to_dict( return data_dict +@dataclass(frozen=True, slots=True) +class ObjectPermissionUpsert: + object_permission_id: str + record: dict[str, object] + + +async def prepare_object_permission_upsert( + new_object_permission: Mapping[str, object], + existing_object_permission_id: str | None, + prisma_client: PrismaClient, +) -> ObjectPermissionUpsert: + """ + Read-and-merge half of an object permission upsert; performs no writes. + + Merges the sent grants over the existing row (looked up by + ``existing_object_permission_id``, or a fresh uuid when the entity has none) and + returns the id plus the full record to upsert. The id is pinned inside the record + because the column has ``@default(uuid())``, so a create without it would mint a + different id than the one the caller links. ``mcp_tool_permissions`` is serialized + to a JSON string to avoid GraphQL parsing issues (e.g. server IDs starting with + "3e64" being interpreted as floats). + + Keeping this separate from the write lets callers run the upsert inside the same + transaction as the row that links ``object_permission_id``, so a rolled-back + update cannot leave permission changes live. + """ + object_permission_id = existing_object_permission_id or str(uuid.uuid4()) + existing_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( + where={"object_permission_id": object_permission_id}, + ) + existing_fields: dict[str, object] = ( + existing_object_permission.model_dump(exclude_unset=True, exclude_none=True) + if existing_object_permission is not None + else {} + ) + merged: dict[str, object] = { + **existing_fields, + **new_object_permission, + "object_permission_id": object_permission_id, + } + record: dict[str, object] = { + **merged, + **( + {"mcp_tool_permissions": safe_dumps(merged["mcp_tool_permissions"])} + if "mcp_tool_permissions" in merged + else {} + ), + } + return ObjectPermissionUpsert(object_permission_id=object_permission_id, record=record) + + async def handle_update_object_permission_common( data_json: Dict, existing_object_permission_id: Optional[str], @@ -93,50 +145,23 @@ async def handle_update_object_permission_common( if prisma_client is None: raise ValueError("Prisma client not found") - ######################################################### - # Ensure `object_permission` is not added to the data_json - # We need to update the entity at the object_permission_id level in the LiteLLM_ObjectPermissionTable - ######################################################### - new_object_permission: Union[dict, str] = data_json.pop("object_permission", None) + new_object_permission: Union[dict, str, None] = data_json.pop("object_permission", None) if new_object_permission is None: return None - # Lookup existing object permission ID and update that entry - object_permission_id_to_use: str = existing_object_permission_id or str(uuid.uuid4()) - existing_object_permissions_dict: Dict = {} - - existing_object_permission = await ObjectPermissionRepository(prisma_client).table.find_unique( - where={"object_permission_id": object_permission_id_to_use}, - ) - - # Update the object permission - if existing_object_permission is not None: - existing_object_permissions_dict = existing_object_permission.model_dump(exclude_unset=True, exclude_none=True) - - # Handle string JSON object permission if isinstance(new_object_permission, str): new_object_permission = json.loads(new_object_permission) - if isinstance(new_object_permission, dict): - existing_object_permissions_dict.update(new_object_permission) - - ######################################################### - # Serialize mcp_tool_permissions JSON field to avoid GraphQL parsing issues - # (e.g., server IDs starting with "3e64" being interpreted as floats) - ######################################################### - if "mcp_tool_permissions" in existing_object_permissions_dict: - existing_object_permissions_dict["mcp_tool_permissions"] = safe_dumps( - existing_object_permissions_dict["mcp_tool_permissions"] - ) - - ######################################################### - # Commit the update to the LiteLLM_ObjectPermissionTable - ######################################################### + upsert = await prepare_object_permission_upsert( + new_object_permission=new_object_permission if isinstance(new_object_permission, dict) else {}, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) created_object_permission_row = await ObjectPermissionRepository(prisma_client).table.upsert( - where={"object_permission_id": object_permission_id_to_use}, + where={"object_permission_id": upsert.object_permission_id}, data={ - "create": existing_object_permissions_dict, - "update": existing_object_permissions_dict, + "create": upsert.record, + "update": upsert.record, }, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index f4470e7e83de..7ed123f6cdfc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -10,9 +10,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../../") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../")) # Adds the parent directory to the system path @pytest.mark.asyncio @@ -58,16 +56,12 @@ async def test_organization_update_object_permissions_existing_permission(monkey "vector_stores": ["old_store_1", "old_store_2"], } - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=existing_object_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_object_permission) # Mock upsert operation updated_permission = MagicMock() updated_permission.object_permission_id = "existing_perm_id_123" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=updated_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=updated_permission) # Test data with new object permission data_json = { @@ -107,9 +101,7 @@ async def test_get_organization_daily_activity_admin_param_passing(monkeypatch): # Mock prisma client mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Admin view -> skip membership restriction @@ -121,9 +113,7 @@ async def test_get_organization_daily_activity_admin_param_passing(monkeypatch): # Patch downstream common function and verify call args mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") get_daily_activity_mock = AsyncMock(return_value=mocked_response) - monkeypatch.setattr( - organization_endpoints, "get_daily_activity", get_daily_activity_mock - ) + monkeypatch.setattr(organization_endpoints, "get_daily_activity", get_daily_activity_mock) auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") result = await get_organization_daily_activity( @@ -172,17 +162,11 @@ async def test_get_organization_daily_activity_non_admin_defaults_to_admin_orgs( # Mock prisma client and memberships mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_organizationmembership.find_many = AsyncMock( return_value=[ - SimpleNamespace( - organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value - ), - SimpleNamespace( - organization_id="orgB", user_role=LitellmUserRoles.ORG_ADMIN.value - ), + SimpleNamespace(organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value), + SimpleNamespace(organization_id="orgB", user_role=LitellmUserRoles.ORG_ADMIN.value), ] ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -196,13 +180,9 @@ async def test_get_organization_daily_activity_non_admin_defaults_to_admin_orgs( # Patch downstream aggregator mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") get_daily_activity_mock = AsyncMock(return_value=mocked_response) - monkeypatch.setattr( - organization_endpoints, "get_daily_activity", get_daily_activity_mock - ) + monkeypatch.setattr(organization_endpoints, "get_daily_activity", get_daily_activity_mock) - auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user" - ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user") await get_organization_daily_activity( organization_ids=None, start_date="2024-02-01", @@ -238,15 +218,9 @@ async def test_get_organization_daily_activity_non_admin_unauthorized_org_raises # Mock prisma client and memberships (only orgA is admin) mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_organizationmembership.find_many = AsyncMock( - return_value=[ - SimpleNamespace( - organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value - ) - ] - ) - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[] + return_value=[SimpleNamespace(organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value)] ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Non-admin view @@ -255,9 +229,7 @@ async def test_get_organization_daily_activity_non_admin_unauthorized_org_raises lambda _: False, ) - auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user" - ) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user") with pytest.raises(HTTPException) as exc: await get_organization_daily_activity( @@ -312,21 +284,17 @@ async def test_organization_update_object_permissions_no_existing_permission( ) # Mock find_unique to return None (no existing permission) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "new_perm_id_456" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=new_permission) data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["brand_new_store"] - ).model_dump(exclude_unset=True, exclude_none=True), + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["brand_new_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), "organization_alias": "updated_org_2", } @@ -381,21 +349,17 @@ async def test_organization_update_object_permissions_missing_permission_record( ) # Mock find_unique to return None (permission record not found) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "recreated_perm_id_789" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(return_value=new_permission) data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["recreated_store"] - ).model_dump(exclude_unset=True, exclude_none=True), + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["recreated_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), "organization_alias": "updated_org_3", } @@ -446,18 +410,14 @@ async def test_list_organization_filter_by_org_id(monkeypatch): ) # Mock find_many to return filtered results - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[mock_org1] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[mock_org1]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Test as proxy admin auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") - result = await list_organization( - org_id="org-123", org_alias=None, user_api_key_dict=auth - ) + result = await list_organization(org_id="org-123", org_alias=None, user_api_key_dict=auth) # Verify the correct organization was returned assert len(result) == 1 @@ -512,18 +472,14 @@ async def test_list_organization_filter_by_org_alias(monkeypatch): ) # Mock find_many to return filtered results - mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( - return_value=[mock_org1, mock_org2] - ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[mock_org1, mock_org2]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Test as proxy admin with org_alias filter auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") - result = await list_organization( - org_id=None, org_alias="test", user_api_key_dict=auth - ) + result = await list_organization(org_id=None, org_alias="test", user_api_key_dict=auth) # Verify organizations with "test" in alias were returned assert len(result) == 2 @@ -532,9 +488,7 @@ async def test_list_organization_filter_by_org_alias(monkeypatch): # Verify find_many was called with correct where conditions (case-insensitive contains) mock_prisma_client.db.litellm_organizationtable.find_many.assert_called_once() call_args = mock_prisma_client.db.litellm_organizationtable.find_many.call_args - assert call_args.kwargs["where"] == { - "organization_alias": {"contains": "test", "mode": "insensitive"} - } + assert call_args.kwargs["where"] == {"organization_alias": {"contains": "test", "mode": "insensitive"}} assert call_args.kwargs["include"] == { "litellm_budget_table": True, "members": True, @@ -612,16 +566,12 @@ def patched_org_prisma(): ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), ): - mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock( - return_value=victim_row - ) + mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock(return_value=victim_row) yield mock_prisma @pytest.mark.asyncio -async def test_organization_member_add_rejects_unauthorized_caller( - patched_org_prisma, unauthorized_caller -): +async def test_organization_member_add_rejects_unauthorized_caller(patched_org_prisma, unauthorized_caller): # ``organization_member_add`` catches HTTPException in its # catch-all and re-wraps as ProxyException with the original status # code preserved. @@ -653,9 +603,7 @@ async def test_organization_member_add_rejects_unauthorized_caller( @pytest.mark.asyncio -async def test_organization_member_update_rejects_unauthorized_caller( - patched_org_prisma, unauthorized_caller -): +async def test_organization_member_update_rejects_unauthorized_caller(patched_org_prisma, unauthorized_caller): from litellm.proxy._types import OrganizationMemberUpdateRequest from litellm.proxy.management_endpoints.organization_endpoints import ( organization_member_update, @@ -676,9 +624,7 @@ async def test_organization_member_update_rejects_unauthorized_caller( @pytest.mark.asyncio -async def test_organization_member_delete_rejects_unauthorized_caller( - patched_org_prisma, unauthorized_caller -): +async def test_organization_member_delete_rejects_unauthorized_caller(patched_org_prisma, unauthorized_caller): from litellm.proxy._types import OrganizationMemberDeleteRequest from litellm.proxy.management_endpoints.organization_endpoints import ( organization_member_delete, @@ -695,3 +641,354 @@ async def test_organization_member_delete_rejects_unauthorized_caller( user_api_key_dict=unauthorized_caller, ) assert exc.value.status_code == 403 + + +@pytest.mark.parametrize( + "body", + [{"tpm_limit": ""}, {"tmp_limit": None}], + ids=["non-numeric-limit", "unknown-key"], +) +def test_v2_model_rejects_invalid_body(body): + """A non-numeric limit and an unknown/misspelled key are both rejected at model validation (422 at the route).""" + from pydantic import ValidationError + + from litellm.proxy._types import OrganizationUpdateRequestV2 + + with pytest.raises(ValidationError): + OrganizationUpdateRequestV2.model_validate(body) + + +class _FakeTxContext: + def __init__(self, tx): + self._tx = tx + + async def __aenter__(self): + return self._tx + + async def __aexit__(self, exc_type, exc, tb): + return False + + +async def _run_update_organization_v2( + monkeypatch, + *, + body: dict, + existing_budget_id, + existing_metadata, + existing_object_permission_id=None, + existing_object_permission_row=None, +): + from litellm.proxy._types import ( + LitellmUserRoles, + OrganizationUpdateRequestV2, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import ( + update_organization_v2, + ) + from litellm.proxy.utils import jsonify_object + + mock_prisma_client = AsyncMock() + mock_prisma_client.jsonify_object = jsonify_object + + existing_org = MagicMock() + existing_org.budget_id = existing_budget_id + existing_org.object_permission_id = existing_object_permission_id + existing_org.metadata = existing_metadata + + mock_prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(return_value=existing_org) + mock_prisma_client.db.litellm_organizationtable.update = AsyncMock(return_value=MagicMock()) + mock_prisma_client.db.litellm_budgettable.update = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=existing_object_permission_row + ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock() + + tx = MagicMock() + tx.litellm_organizationtable = mock_prisma_client.db.litellm_organizationtable + tx.litellm_budgettable = mock_prisma_client.db.litellm_budgettable + tx.litellm_objectpermissiontable.upsert = AsyncMock() + mock_prisma_client.db.tx = MagicMock(return_value=_FakeTxContext(tx)) + mock_prisma_client.tx = tx + + call_order = MagicMock() + call_order.attach_mock(tx.litellm_objectpermissiontable.upsert, "permission_upsert") + call_order.attach_mock(mock_prisma_client.db.litellm_organizationtable.update, "org_update") + mock_prisma_client.call_order = call_order + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_verify_org_access", AsyncMock()) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate(body), + user_api_key_dict=auth, + ) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_v2_update_clears_tpm_limit_and_metadata(monkeypatch): + """A cleared tpm_limit is written to the budget row as None; a cleared metadata is written as {}.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"tpm_limit": None, "metadata": None}, + existing_budget_id="budget-1", + existing_metadata={"stale": "value"}, + ) + + budget_write = prisma.db.litellm_budgettable.update.await_args + assert budget_write.kwargs["where"] == {"budget_id": "budget-1"} + assert budget_write.kwargs["data"]["tpm_limit"] is None + assert "soft_budget" not in budget_write.kwargs["data"] + + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert json.loads(write_data["metadata"]) == {} + assert "budget_id" not in write_data + + +@pytest.mark.asyncio +async def test_v2_update_untouched_fields_not_written(monkeypatch): + """Omitted fields are left untouched: only organization_alias is written, no budget-row write.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"organization_alias": "renamed"}, + existing_budget_id="budget-1", + existing_metadata={"keep": "me"}, + ) + + prisma.db.litellm_budgettable.update.assert_not_awaited() + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["organization_alias"] == "renamed" + assert "metadata" not in write_data + assert "tpm_limit" not in write_data + + +@pytest.mark.asyncio +async def test_v2_update_metadata_replaces_not_merges(monkeypatch): + """Sending metadata replaces the stored blob wholesale; a previously-present key is gone.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"metadata": {"a": 1}}, + existing_budget_id="budget-1", + existing_metadata={"stale": "value"}, + ) + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert json.loads(write_data["metadata"]) == {"a": 1} + + +@pytest.mark.asyncio +async def test_v2_rejects_null_clear_of_non_nullable_fields(monkeypatch): + """organization_alias and models are non-nullable columns, so a null clear is a 422, not a 500.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + + for body in ({"organization_alias": None}, {"models": None}): + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate(body), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 422 + + +@pytest.mark.asyncio +async def test_v2_rejects_negative_max_budget(monkeypatch): + """v2 rejects a negative max_budget with a 422 before touching the DB.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate({"max_budget": -5}), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 422 + assert "max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_v2_rejects_caller_without_org_access(monkeypatch): + """v2 runs the real _verify_org_access guard: a non-admin without ORG_ADMIN on the org gets 403 and no write.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_user_has_admin_view", lambda _: False) + + caller = MagicMock() + caller.organization_memberships = [] + monkeypatch.setattr(organization_endpoints, "get_user_object", AsyncMock(return_value=caller)) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1") + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate({"tpm_limit": 5}), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_organizationtable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v2_wires_object_permission_onto_org_write(monkeypatch): + """A sent object_permission merges over the existing permission row and its id is linked onto the org write.""" + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "object_permission_id": "op-123", + "mcp_servers": ["server-1"], + } + + prisma = await _run_update_organization_v2( + monkeypatch, + body={"object_permission": {"vector_stores": ["vs-1"]}}, + existing_budget_id="budget-1", + existing_metadata={}, + existing_object_permission_id="op-123", + existing_object_permission_row=existing_row, + ) + + upsert = prisma.tx.litellm_objectpermissiontable.upsert.await_args.kwargs + assert upsert["where"] == {"object_permission_id": "op-123"} + assert upsert["data"]["update"]["mcp_servers"] == ["server-1"] + assert upsert["data"]["update"]["vector_stores"] == ["vs-1"] + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["object_permission_id"] == "op-123" + + +@pytest.mark.asyncio +async def test_v2_object_permission_upsert_runs_inside_transaction(monkeypatch): + """The permission upsert runs on the tx client, before the org write that links it, so a rollback cannot + leave merged grants live on a row the org still points at.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"object_permission": {"vector_stores": ["vs-1"]}}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + prisma.tx.litellm_objectpermissiontable.upsert.assert_awaited_once() + prisma.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + + upsert = prisma.tx.litellm_objectpermissiontable.upsert.await_args.kwargs + linked_id = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"]["object_permission_id"] + assert upsert["where"] == {"object_permission_id": linked_id} + assert upsert["data"]["create"]["object_permission_id"] == linked_id + + ordered = [name for name, _, _ in prisma.call_order.mock_calls if name in ("permission_upsert", "org_update")] + assert ordered == ["permission_upsert", "org_update"] + + +@pytest.mark.asyncio +async def test_v2_clears_object_permission_when_sent_null(monkeypatch): + """object_permission: null detaches the org's permission row (object_permission_id -> None), no merge.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"object_permission": None}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + prisma.tx.litellm_objectpermissiontable.upsert.assert_not_awaited() + prisma.db.litellm_objectpermissiontable.find_unique.assert_not_awaited() + write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] + assert write_data["object_permission_id"] is None + + +@pytest.mark.asyncio +async def test_v2_rejects_empty_object_permission(monkeypatch): + """object_permission: {} merges nothing, so it is rejected (send null to clear) rather than silently leaving grants.""" + from litellm.proxy._types import LitellmUserRoles, OrganizationUpdateRequestV2, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import update_organization_v2 + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr(organization_endpoints, "_verify_org_access", AsyncMock()) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1") + with pytest.raises(HTTPException) as exc: + await update_organization_v2( + organization_id="org-1", + data=OrganizationUpdateRequestV2.model_validate({"object_permission": {}}), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 422 + assert "object_permission" in str(exc.value.detail) + mock_prisma_client.db.litellm_organizationtable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_v2_writes_budget_and_org_in_one_transaction(monkeypatch): + """A change touching both the budget row and the org row runs both writes inside one prisma transaction.""" + prisma = await _run_update_organization_v2( + monkeypatch, + body={"tpm_limit": 500, "metadata": {"a": 1}}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + prisma.db.tx.assert_called_once() + prisma.db.litellm_budgettable.update.assert_awaited_once() + prisma.db.litellm_organizationtable.update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_v2_serializes_model_max_budget_on_budget_write(monkeypatch): + """model_max_budget is a Json column, so it is JSON-serialized on the budget-row write like new_budget/metadata.""" + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_model_max_budget", + lambda _: None, + ) + + prisma = await _run_update_organization_v2( + monkeypatch, + body={"model_max_budget": {"gpt-4o": {"max_budget": 10}}}, + existing_budget_id="budget-1", + existing_metadata={}, + ) + + written = prisma.db.litellm_budgettable.update.await_args.kwargs["data"]["model_max_budget"] + assert isinstance(written, str) + assert json.loads(written) == {"gpt-4o": {"max_budget": 10}} + + +def test_build_budget_write_data_recomputes_reset_at_on_duration(): + """A sent budget_duration recomputes budget_reset_at so the reset window follows the new duration.""" + from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data + + data = build_budget_write_data({"budget_duration": "30d"}, "admin-1") + assert data["budget_duration"] == "30d" + assert "budget_reset_at" in data + assert data["updated_by"] == "admin-1" + + +def test_build_budget_write_data_no_reset_at_without_duration(): + """Clearing a limit writes it through untouched and does not recompute budget_reset_at.""" + from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data + + data = build_budget_write_data({"tpm_limit": None}, "admin-1") + assert data["tpm_limit"] is None + assert "budget_reset_at" not in data + + +def test_build_budget_write_data_clears_reset_at_with_null_duration(): + """Clearing budget_duration also nulls budget_reset_at so no stale reset timestamp survives.""" + from litellm.proxy.management_endpoints.organization_endpoints import build_budget_write_data + + data = build_budget_write_data({"budget_duration": None}, "admin-1") + assert data["budget_duration"] is None + assert data["budget_reset_at"] is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e612bddaca17..e0223aa4af68 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -18887,6 +18887,33 @@ export interface paths { patch?: never; trace?: never; }; + "/v2/organization/{organization_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + /** + * Update Organization V2 + * @description Partial update of an organization (RESTful PATCH, RFC 7396 merge-patch semantics). + * + * A sent field is written and an omitted one is left untouched (presence is read from + * ``model_fields_set``). Clear tokens are per field: budget limits and ``metadata`` clear with + * ``null``, ``models`` with ``[]``, and ``object_permission`` with ``null`` (it merges when sent, + * so an empty ``{}`` is rejected). ``organization_alias`` is required and cannot be cleared. + * Validation failures return 422; the object-permission upsert, budget-row write, and + * org-row write are one transaction. + */ + patch: operations["update_organization_v2_v2_organization__organization_id__patch"]; + trace?: never; + }; "/v2/rerank": { parameters: { query?: never; @@ -28642,6 +28669,41 @@ export interface components { /** Organizations */ organizations: string[]; }; + /** + * OrganizationUpdateRequestV2 + * @description Typed PATCH body for ``/v2/organization/{organization_id}`` (RFC 7396 merge-patch). + * + * Presence is read from ``model_fields_set``, so a sent field is written and an omitted one is + * left untouched. ``extra="forbid"`` makes an unknown key a 422 rather than a silent no-op, since + * the contract hinges on which keys are present. See the endpoint for the per-field clear tokens. + */ + OrganizationUpdateRequestV2: { + /** Budget Duration */ + budget_duration?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Max Parallel Requests */ + max_parallel_requests?: number | null; + /** Metadata */ + metadata?: { + [key: string]: unknown; + } | null; + /** Model Max Budget */ + model_max_budget?: { + [key: string]: unknown; + } | null; + /** Models */ + models?: string[] | null; + object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null; + /** Organization Alias */ + organization_alias?: string | null; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + /** Tpm Limit */ + tpm_limit?: number | null; + }; /** * PaginatedAuditLogResponse * @description Response model for paginated audit logs @@ -57644,6 +57706,41 @@ export interface operations { }; }; }; + update_organization_v2_v2_organization__organization_id__patch: { + parameters: { + query?: never; + header?: never; + path: { + organization_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["OrganizationUpdateRequestV2"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_OrganizationTableWithMembers"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; rerank_v2_rerank_post: { parameters: { query?: never;