diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 154a60f1cd1..b17498a370a 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -26,6 +26,7 @@ from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import delete_cached_project_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.guardrails.guardrail_helpers import validate_guardrail_names from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamAccess, is_team_admin, is_team_member from litellm.proxy.management.teams.dependencies import get_team_access from litellm.proxy.management.users.service import org_admin_org_ids @@ -604,6 +605,13 @@ async def new_project( "error": f"Only admins or team admins can create projects. Your role is {user_api_key_dict.user_role}" }, ) + await validate_guardrail_names( + None, + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + None, + prisma_client, + license_first=False, + ) quota_policy: Final = _project_model_quota_policy(general_settings_view()) if quota_policy is not None: @@ -835,6 +843,13 @@ async def update_project( status_code=403, detail={"error": "Cannot reassign project to a team you are not an admin of"}, ) + await validate_guardrail_names( + None, + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + existing_project.metadata if isinstance(existing_project.metadata, dict) else None, + prisma_client, + license_first=False, + ) # Validate project limits against team limits if target_team_obj is not None: diff --git a/litellm/proxy/guardrails/guardrail_helpers.py b/litellm/proxy/guardrails/guardrail_helpers.py index 3282334715d..eb258ae7de7 100644 --- a/litellm/proxy/guardrails/guardrail_helpers.py +++ b/litellm/proxy/guardrails/guardrail_helpers.py @@ -1,13 +1,21 @@ -import os -import sys -from typing import Final +import json +from collections.abc import Mapping, Sequence +from itertools import chain +from typing import TYPE_CHECKING, Final, Literal + +from fastapi import HTTPException +from pydantic import StrictBool, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy.proxy_server import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.utils import premium_user_check +from litellm.repositories.table_repositories import GuardrailsRepository from litellm.types.guardrails import * -sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient def can_modify_guardrails(team_obj: LiteLLM_TeamTable | None) -> bool: @@ -104,3 +112,79 @@ async def should_proceed_based_on_api_key(user_api_key_dict: UserAPIKeyAuth, gua # Do not proceeed if - "metadata": { "guardrails": { "lakera_prompt_injection": false } } return False return True + + +_GUARDRAIL_ENTRIES: Final = TypeAdapter(list[str | dict[str, object]]) +_MODIFY_GUARDRAILS_PERMISSION: Final = TypeAdapter(dict[Literal["modify_guardrails"], StrictBool]) + + +def _guardrail_entry_names(entries: object) -> frozenset[str] | None: + try: + parsed: Final = _GUARDRAIL_ENTRIES.validate_python(entries) + except ValidationError: + return None + return frozenset(chain.from_iterable((entry,) if isinstance(entry, str) else entry.keys() for entry in parsed)) + + +def _is_modify_guardrails_permission(entries: object) -> bool: + # a team's documented {"modify_guardrails": false}, read by auth checks, is the one stored dict that has an effect + try: + return len(_MODIFY_GUARDRAILS_PERMISSION.validate_python(entries, strict=True)) == 1 + except ValidationError: + return False + + +def _malformed_guardrails_message(entries: object) -> str: + shown: Final = json.dumps(entries, default=str) + return ( + f"Invalid guardrail detected - {shown if len(shown) <= 100 else shown[:100] + '...'}. " + + "Please provide a valid list of guardrail names." + ) + + +async def registered_guardrail_names(names: frozenset[str], prisma_client: "PrismaClient | None") -> frozenset[str]: + loaded_names: Final = frozenset( + callback.guardrail_name + for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + if isinstance(callback, CustomGuardrail) + ) + loaded: Final = names & loaded_names + missing: Final = names - loaded + if not missing or prisma_client is None: + return loaded + return loaded | await GuardrailsRepository(prisma_client, use_writer=True).find_active_guardrail_names(missing) + + +async def validate_guardrail_names( + guardrails: Sequence[str] | None, + metadata: Mapping[str, object] | None, + existing_metadata: Mapping[str, object] | None, + prisma_client: "PrismaClient | None", + *, + license_first: bool, +) -> None: + if license_first and guardrails: + # key and team writes license-gate top-level guardrails later on, keep that 403 ahead of a name 400 + premium_user_check("guardrails") + requested_entries: Final = ( + guardrails if guardrails is not None else (metadata.get("guardrails") if metadata is not None else None) + ) + if requested_entries is None or _is_modify_guardrails_permission(requested_entries): + return + requested: Final = _guardrail_entry_names(requested_entries) + if requested is None: + raise HTTPException(status_code=400, detail={"error": _malformed_guardrails_message(requested_entries)}) + stored_entries: Final = existing_metadata.get("guardrails") if existing_metadata is not None else None + new_names: Final = requested - (_guardrail_entry_names(stored_entries) or frozenset()) + if not new_names: + return + unknown: Final = new_names - await registered_guardrail_names(new_names, prisma_client) + if not unknown: + return + raise HTTPException( + status_code=400, + detail={ + "error": f"Unknown guardrail(s): {', '.join(sorted(unknown))}. " + + "Create the guardrail first, or check the name against /v2/guardrails/list." + }, + ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 3307f3050d9..f973fadb53c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -58,6 +58,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( user_object_permission_id_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.guardrails.guardrail_helpers import validate_guardrail_names from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks @@ -601,6 +602,13 @@ async def new_user( enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key") enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key") raise_on_logging_credential_change(data.metadata, None, user_api_key_dict) + await validate_guardrail_names( + data.guardrails, + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + None, + prisma_client, + license_first=False, + ) check_permissions_caller_permission( data=data, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 90779f732c7..3bd17c4b5d8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -86,6 +86,7 @@ from litellm.proxy.common_utils.openai_error_payload import openai_error_param, from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key +from litellm.proxy.guardrails.guardrail_helpers import validate_guardrail_names from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management.teams.authz import TEAM_ADMIN_ONLY, TEAM_OR_ORG_ADMIN, is_team_admin @@ -1615,6 +1616,7 @@ async def _common_key_generation_helper( _requested_metadata, user_api_key_dict, ) + await validate_guardrail_names(data.guardrails, _requested_metadata, None, prisma_client, license_first=True) # APPLY ENTERPRISE KEY MANAGEMENT PARAMS try: @@ -3403,6 +3405,13 @@ async def _process_single_key_update( user_api_key_dict, existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) + await validate_guardrail_names( + update_key_request.guardrails, + update_key_request.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict + prisma_client, + license_first=True, + ) # Custom key update hook if user_custom_key_update is not None: @@ -3957,6 +3966,13 @@ async def _validate_update_key_data( user_api_key_dict, existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) + await validate_guardrail_names( + data.guardrails, + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict + checked_prisma_client, + license_first=True, + ) # Check team limits if key has a team_id (from request or existing key) team_obj: LiteLLM_TeamTableCachedObj | None = None @@ -6752,6 +6768,13 @@ async def regenerate_key_fn( user_api_key_dict, existing_metadata=_key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict ) + await validate_guardrail_names( + data.guardrails if data is not None else None, + data.metadata if data is not None else None, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + _key_in_db.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict + prisma_client, + license_first=True, + ) if data is not None and (data.access_group_ids or data.object_permission is not None): regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0704d2c2444..f3d3c01b3a4 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -142,6 +142,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_reservation_cache_key, ) from litellm.proxy.db.db_span import db_span +from litellm.proxy.guardrails.guardrail_helpers import validate_guardrail_names from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( RoutedModels, @@ -1917,6 +1918,13 @@ async def new_team( user_api_key_dict, metadata_sent="metadata" in data.model_fields_set, ) + await validate_guardrail_names( + data.guardrails, + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + None, + prisma_client, + license_first=True, + ) if isinstance(data.metadata, dict): TeamMemberBudgetHandler.strip_system_managed_metadata_keys(data.metadata) @@ -2637,6 +2645,13 @@ async def update_team( metadata_sent="metadata" in data.model_fields_set, existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # existing_team_row.metadata is a bare dict ) + await validate_guardrail_names( + data.guardrails, + data.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # request models declare `metadata` as bare dict + _existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # existing_team_row.metadata is a bare dict + prisma_client, + license_first=True, + ) if data.soft_budget is not None: max_budget_to_check = data.max_budget if data.max_budget is not None else existing_team_row.max_budget diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 3c3b87c9b8c..99d2ef97634 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction +from litellm.repositories.chunked_in import find_many_in from litellm.repositories.prisma_protocols import RowT_co, TableActions if TYPE_CHECKING: @@ -103,6 +104,10 @@ class ObjectPermissionRepository(PrismaTableRepository["prisma_models.LiteLLM_Ob class GuardrailsRepository(PrismaTableRepository["prisma_models.LiteLLM_GuardrailsTable"]): table_name = "litellm_guardrailstable" + async def find_active_guardrail_names(self, names: frozenset[str]) -> frozenset[str]: + rows: Final = await find_many_in(self.table, "guardrail_name", sorted(names), where={"status": "active"}) + return frozenset(row.guardrail_name for row in rows) + class MCPServerRepository(PrismaTableRepository["prisma_models.LiteLLM_MCPServerTable"]): table_name = "litellm_mcpservertable" diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 2a37ed2ed1f..95fe251564a 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -2363,3 +2363,51 @@ async def test_project_info_missing_project_is_404_even_for_org_admins(project_a ) assert str(exc_info.value.code) == "404" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ("create", "update")) +async def test_project_writes_reject_an_unregistered_guardrail(monkeypatch, operation): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import ProxyException + + team_row = MagicMock() + team_row.model_dump = lambda: {"team_id": "team-1"} + existing_row = MagicMock(team_id="team-1", budget_id=None, object_permission_id=None) + existing_row.metadata = {"guardrails": ["stale"]} + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.create = AsyncMock() + mock_prisma.db.litellm_projecttable.update = AsyncMock() + mock_prisma.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key=MASTER_KEY, user_id="1234") + + write: Final = ( + new_project( + data=NewProjectRequest(team_id="team-1", guardrails=["typo"]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=admin, + ) + if operation == "create" + else update_project( + data=UpdateProjectRequest(project_id="project-1", guardrails=["stale", "typo"]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=admin, + ) + ) + with pytest.raises((HTTPException, ProxyException)) as exc: + await write + + message: Final = str(getattr(exc.value, "message", None) or getattr(exc.value, "detail", "")) + assert "typo" in message + assert "stale" not in message + assert str(getattr(exc.value, "code", None) or getattr(exc.value, "status_code", "")) == "400" + mock_prisma.db.litellm_projecttable.create.assert_not_awaited() + mock_prisma.db.litellm_projecttable.update.assert_not_awaited() diff --git a/tests/unit/proxy/guardrails/test_guardrail_helpers.py b/tests/unit/proxy/guardrails/test_guardrail_helpers.py new file mode 100644 index 00000000000..d30918d4b2b --- /dev/null +++ b/tests/unit/proxy/guardrails/test_guardrail_helpers.py @@ -0,0 +1,228 @@ +from collections.abc import Iterator, Mapping +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.guardrail_helpers import validate_guardrail_names + +DB_ROWS: Final = MappingProxyType({"valid": "active", "also-valid": "active", "submitted": "pending_review"}) + + +def _filter(where: Mapping[str, object]) -> Mapping[str, object]: + clauses: Final = where.get("AND", (where,)) + return {key: value for clause in clauses for key, value in clause.items()} # comprehension-ok: flattens prisma AND clauses + + +def _prisma(rows: Mapping[str, str] = DB_ROWS) -> MagicMock: + async def find_many(*, where: Mapping[str, object]) -> list[SimpleNamespace]: + flat: Final = _filter(where) + names: Final = flat["guardrail_name"]["in"] + return [SimpleNamespace(guardrail_name=name) for name in names if rows.get(name) == flat.get("status")] + + client: Final = MagicMock() + client.writer_db.litellm_guardrailstable.find_many = AsyncMock(side_effect=find_many) + client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + return client + + +def _looked_up(client: MagicMock) -> list[frozenset[str]]: + calls: Final = client.writer_db.litellm_guardrailstable.find_many.await_args_list + return [frozenset(_filter(call.kwargs["where"])["guardrail_name"]["in"]) for call in calls] + + +async def _validate( + guardrails: list[str] | None, + metadata: Mapping[str, object] | None, + client: MagicMock, + existing_metadata: Mapping[str, object] | None = None, +) -> None: + await validate_guardrail_names(guardrails, metadata, existing_metadata, client, license_first=False) + + +@pytest.fixture +def guardrail_loaded_here() -> Iterator[str]: + lists: Final = [ + litellm.callbacks, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ] + snapshots: Final = [list(cb_list) for cb_list in lists] + litellm.logging_callback_manager.add_litellm_callback(CustomGuardrail(guardrail_name="loaded-here")) + try: + yield "loaded-here" + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guardrails", "metadata"), + ( + (["valid", "typo"], None), + (None, {"guardrails": ["valid", "typo"]}), + (None, {"guardrails": [{"valid": {}}, {"typo": {"extra_body": {"k": "v"}}}]}), + ), + ids=("top_level", "metadata_names", "metadata_dynamic_params"), +) +async def test_rejects_only_the_unregistered_name( + guardrails: list[str] | None, metadata: dict[str, object] | None +) -> None: + with pytest.raises(HTTPException) as exc: + await _validate(guardrails, metadata, _prisma()) + + assert exc.value.status_code == 400 + message: Final = exc.value.detail["error"] + assert "typo" in message + assert "valid" not in message + + +@pytest.mark.asyncio +async def test_active_db_rows_and_loaded_guardrails_pass(guardrail_loaded_here: str) -> None: + client: Final = _prisma() + + await _validate(None, {"guardrails": ["valid", {"also-valid": {}}, guardrail_loaded_here]}, client) + + assert _looked_up(client) == [frozenset({"valid", "also-valid"})] + + +@pytest.mark.asyncio +async def test_a_guardrail_pending_review_is_unknown() -> None: + with pytest.raises(HTTPException) as exc: + await _validate(["submitted"], None, _prisma()) + + assert "submitted" in exc.value.detail["error"] + + +@pytest.mark.asyncio +async def test_the_lookup_reads_the_writer_so_a_lagging_replica_cannot_hide_a_new_guardrail() -> None: + client: Final = _prisma() + + await _validate(["valid"], None, client) + + client.db.litellm_guardrailstable.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_no_db_checks_only_loaded_guardrails(guardrail_loaded_here: str) -> None: + await validate_guardrail_names([guardrail_loaded_here], None, None, None, license_first=False) + with pytest.raises(HTTPException): + await validate_guardrail_names(["valid"], None, None, None, license_first=False) + + +@pytest.mark.asyncio +async def test_top_level_guardrails_replace_the_metadata_list() -> None: + client: Final = _prisma() + + await _validate(["valid"], {"guardrails": ["typo"]}, client) + + assert _looked_up(client) == [frozenset({"valid"})] + + +@pytest.mark.asyncio +async def test_names_already_on_the_object_are_not_rechecked() -> None: + client: Final = _prisma() + stored: Final = {"guardrails": ["stale"]} + + await _validate(["stale"], None, client, stored) + await _validate(["stale", "valid"], None, client, stored) + with pytest.raises(HTTPException): + await _validate(["stale", "typo"], None, client, stored) + + assert _looked_up(client) == [frozenset({"valid"}), frozenset({"typo"})] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guardrails", "metadata"), + ((None, None), (None, {"tags": ["t"]}), (None, {"guardrails": None}), ([], None), (None, {"guardrails": []})), + ids=("nothing", "other_metadata", "metadata_null", "empty_top_level", "empty_metadata"), +) +async def test_no_requested_names_skips_the_lookup( + guardrails: list[str] | None, metadata: dict[str, object] | None +) -> None: + client: Final = _prisma() + + await _validate(guardrails, metadata, client) + + assert _looked_up(client) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entries", "shown"), + (("pii", '"pii"'), (5, "5"), ([1], "[1]"), ([None, "stale"], '[null, "stale"]')), + ids=("string", "number", "non_name_entry", "null_entry"), +) +async def test_malformed_metadata_guardrails_fail_and_name_the_value(entries: object, shown: str) -> None: + with pytest.raises(HTTPException) as exc: + await _validate(None, {"guardrails": entries}, _prisma()) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"] == ( + f"Invalid guardrail detected - {shown}. Please provide a valid list of guardrail names." + ) + + +@pytest.mark.asyncio +async def test_resending_a_malformed_stored_value_still_fails() -> None: + with pytest.raises(HTTPException) as exc: + await _validate(None, {"guardrails": "pii", "note": "unrelated edit"}, _prisma(), {"guardrails": "pii"}) + + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_malformed_message_quotes_a_long_value_only_in_part() -> None: + with pytest.raises(HTTPException) as exc: + await _validate(None, {"guardrails": "g" * 5000}, _prisma()) + + message: Final = exc.value.detail["error"] + assert '"' + "g" * 99 + "..." in message + assert len(message) < 200 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("license_first", "status_code"), ((True, 403), (False, 400))) +async def test_license_first_puts_the_license_error_ahead_of_an_unknown_name( + monkeypatch: pytest.MonkeyPatch, license_first: bool, status_code: int +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + + with pytest.raises(HTTPException) as exc: + await validate_guardrail_names(["typo"], None, None, _prisma(), license_first=license_first) + + assert exc.value.status_code == status_code + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed", (False, True)) +async def test_the_team_modify_guardrails_permission_is_left_alone(allowed: bool) -> None: + client: Final = _prisma() + permission: Final = {"modify_guardrails": allowed} + + await _validate(None, {"guardrails": permission}, client, {"guardrails": permission}) + await _validate(None, {"guardrails": permission}, client) + + assert _looked_up(client) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entries", + ({"valid": True}, {"modify_guardrails": False, "valid": True}, {"modify_guardrails": "no"}), + ids=("toggle_map", "permission_plus_attachment", "non_bool_permission"), +) +async def test_other_stored_dicts_fail_loudly_since_they_never_run(entries: dict[str, object]) -> None: + with pytest.raises(HTTPException) as exc: + await _validate(None, {"guardrails": entries}, _prisma()) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"].startswith("Invalid guardrail detected - {") diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 0ebd6a3ff3a..c67cb974583 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -14265,6 +14265,7 @@ async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): def _policy_generate_prisma() -> MagicMock: mock_prisma = MagicMock() mock_prisma.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1")) + mock_prisma.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[SimpleNamespace(guardrail_name="g1")]) mock_prisma.jsonify_object = MagicMock(side_effect=lambda data: json.loads(data) if isinstance(data, str) else data) return mock_prisma @@ -24464,3 +24465,156 @@ def test_non_admin_resaving_referencing_key_metadata_read_back_encrypted_is_not_ raise_on_logging_credential_change({"note": "unrelated edit"}, key_row, internal_user) assert removed.value.status_code == 403 + + +@pytest.mark.parametrize( + "body", + ({"guardrails": ["typo"]}, {"metadata": {"guardrails": ["typo"]}}), + ids=("top_level", "metadata"), +) +def test_key_generate_rejects_an_unregistered_guardrail( + spend_permission_client: TestClient, monkeypatch: pytest.MonkeyPatch, body: dict[str, object] +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + prisma_client: Final = proxy_server.prisma_client + prisma_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + response: Final = spend_permission_client.post("/key/generate", json=body) + + assert response.status_code == 400, response.text + assert "typo" in response.json()["error"]["message"] + prisma_client.insert_data.assert_not_awaited() + + +def test_key_generate_accepts_a_guardrail_active_in_the_db( + spend_permission_client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "premium_user", True) + prisma_client: Final = proxy_server.prisma_client + prisma_client.writer_db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[SimpleNamespace(guardrail_name="created-on-other-pod")] + ) + + response: Final = spend_permission_client.post("/key/generate", json={"guardrails": ["created-on-other-pod"]}) + + assert response.status_code == 200, response.text + prisma_client.insert_data.assert_awaited_once() + + +def _key_with_stale_guardrail() -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token=_POLICY_HASHED_TOKEN, + user_id="test-user", + team_id="team-a", + metadata={"guardrails": ["stale"]}, + ) + + +@pytest.mark.asyncio +async def test_update_key_fn_rejects_an_added_unregistered_guardrail_and_keeps_stale_keys_editable(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client: Final = _setup_update_key_fn_policy_mocks(monkeypatch, _key_with_stale_guardrail()) + mock_prisma_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + with patch( # test-quality-ok: cache eviction is outside the guardrail check + "litellm.proxy.management_endpoints.key_management_endpoints.delete_cache_key_object", + new_callable=AsyncMock, + ): + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, guardrails=["stale", "typo"]), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + assert str(exc_info.value.code) == "400" + assert "typo" in exc_info.value.message + assert "stale" not in exc_info.value.message + mock_prisma_client.update_data.assert_not_awaited() + + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, guardrails=["stale"], max_budget=50.0), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + mock_prisma_client.update_data.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_process_single_key_update_rejects_an_unregistered_guardrail(monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + mock_prisma_client: Final = AsyncMock() + mock_prisma_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy( + mock_prisma_client, UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, metadata={"guardrails": ["typo"]}), None + ) + + assert exc_info.value.status_code == 400 + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_regenerate_key_rejects_an_unregistered_guardrail(): + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import regenerate_key_fn + + mock_prisma_client: Final = AsyncMock() + mock_prisma_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + mock_repo: Final = MagicMock() + mock_repo.table.find_unique = AsyncMock(return_value=_make_regenerate_existing_key()) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: guardrails are an enterprise field + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: fake DB + patch( # test-quality-ok: the stored key row is outside the guardrail check + "litellm.proxy.management_endpoints.key_management_endpoints.VerificationTokenRepository", + return_value=mock_repo, + ), + pytest.raises(ProxyException) as exc, + ): + await regenerate_key_fn( + key=None, + data=RegenerateKeyRequest(key="sk-alice-personal", guardrails=["typo"]), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert int(exc.value.code) == 400 + assert "typo" in str(exc.value.message) + mock_repo.table.update.assert_not_called() + + +def test_user_new_rejects_an_unregistered_guardrail_before_creating_anything( + spend_permission_client: TestClient, +) -> None: + from litellm.proxy import proxy_server + + prisma_client: Final = proxy_server.prisma_client + prisma_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + response: Final = spend_permission_client.post("/user/new", json={"guardrails": ["typo"]}) + + assert response.status_code == 400, response.text + assert "typo" in response.json()["error"]["message"] + prisma_client.insert_data.assert_not_awaited() + + +def test_user_new_without_a_key_does_not_check_guardrail_names(spend_permission_client: TestClient) -> None: + from litellm.proxy import proxy_server + + find_many: Final = AsyncMock(return_value=[]) + proxy_server.prisma_client.writer_db.litellm_guardrailstable.find_many = find_many + + response: Final = spend_permission_client.post("/user/new", json={"guardrails": ["typo"], "auto_create_key": False}) + + assert response.status_code == 200, response.text + find_many.assert_not_awaited() diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 29e9be69776..58a852247e7 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -9246,6 +9246,12 @@ async def test_update_team_guardrails_with_org_id( mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) + mock_prisma.writer_db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[ + SimpleNamespace(guardrail_name="aporia-pre-call"), + SimpleNamespace(guardrail_name="aporia-post-call"), + ] + ) mock_cache.async_set_cache = AsyncMock() # Mock organization fetch - this is where the bug occurred @@ -19573,3 +19579,62 @@ async def test_team_info_keeps_the_auth_label_for_a_caller_outside_the_team(monk ) assert (exc_info.value.code, exc_info.value.type) == ("403", "auth_error") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fields", + ({"guardrails": ["typo"]}, {"metadata": {"guardrails": ["typo"]}}), + ids=("top_level", "metadata"), +) +async def test_new_team_rejects_an_unregistered_guardrail_before_writing( + mock_db_client, mock_admin_auth, monkeypatch, fields +): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.create = AsyncMock() + mock_db_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + with pytest.raises(ProxyException) as exc: + await new_team( + data=NewTeamRequest(team_alias="gr-team", **fields), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert str(exc.value.code) == "400" + assert "typo" in str(exc.value.message) + mock_db_client.db.litellm_teamtable.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_team_rejects_an_added_unregistered_guardrail(mock_db_client, mock_admin_auth, monkeypatch): + from fastapi import Request + + from litellm.proxy._types import ProxyException, UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + existing = MagicMock() + existing.metadata = {"guardrails": ["stale"]} + existing.model_dump.return_value = {"team_id": "t1", "metadata": {"guardrails": ["stale"]}} + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) + mock_db_client.db.litellm_teamtable.update = AsyncMock() + mock_db_client.writer_db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="t1", guardrails=["stale", "typo"]), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert str(exc.value.code) == "400" + assert "typo" in str(exc.value.message) + assert "stale" not in str(exc.value.message) + mock_db_client.db.litellm_teamtable.update.assert_not_awaited()