mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): reject unknown guardrail names on key, user, team and project writes (#51080)
* fix(proxy): reject unknown guardrail names on key and user writes
Key generate, update, regenerate, team bulk update and /user/new now return 400 for a guardrail name that no loaded guardrail answers to and that has no active DB row, instead of storing it so requests run unguarded. Names the key already holds are not rechecked, and the key endpoints keep the enterprise license 403 ahead of the name check.
Resolves LIT-10613
* fix(proxy): read guardrail names from the writer when validating key writes
A guardrail created on another pod may not have reached a lagging read replica yet, so the fallback lookup now reads the writer and a just-created guardrail is not rejected as unknown. Other guardrail table reads keep using the reader.
* feat(proxy): validate guardrail names on team and project writes, reject malformed values loudly
Team and project create and update now return 400 for guardrail names that no loaded guardrail answers to and that have no active DB row, the same check key and user writes use. A metadata.guardrails value that is not a list of names is rejected even when it equals the stored value, with a short message naming the value. The lookup and the license-first wrapper move into common_utils so every entity shares them.
* refactor(proxy): one validate_guardrail_names entry point for guardrail name checks
The name check, its DB lookup and the license-first option now live in litellm/proxy/guardrails/guardrail_name_validation.py and every key, user, team and project write calls validate_guardrail_names directly. The lookup factory, both wrappers and the per-endpoint helpers are gone, and common_utils and guardrail_registry are back to their main versions. Behavior is unchanged.
* fix(proxy): chunk the guardrail lookup and leave legacy dict guardrail settings alone
The guardrail name lookup goes through find_many_in so its IN list has a written bound, and the project test keeps pytest.raises to one statement. A dict-shaped metadata.guardrails, a team's documented {"modify_guardrails": false} or the v1 {name: bool} toggle map, is read at request time, so validation leaves it alone instead of rejecting it.
* fix(proxy): only let the team modify_guardrails permission through as a dict
Stored key, team and project guardrails only run when metadata.guardrails is a list, so a dict such as {"config-guard": true} saved fine and attached nothing. The one stored dict with an effect is a team's documented {"modify_guardrails": bool}, which auth checks read, so that exact shape passes and every other dict gets the Invalid guardrail 400.
* refactor(proxy): move guardrail name validation into guardrail_helpers
validate_guardrail_names and its helpers now live in litellm/proxy/guardrails/guardrail_helpers.py next to the other guardrail helpers, and guardrail_name_validation.py is removed. guardrail_helpers imports LiteLLM_TeamTable and UserAPIKeyAuth from litellm.proxy._types instead of litellm.proxy.proxy_server, so the endpoints can import it at module level without pulling in the proxy server, and its dead sys.path line is dropped. Behavior is unchanged.
This commit is contained in:
mrinal-berri
GitHub
parent
7461a996cc
commit
c96ca27afc
10 changed files with 650 additions and 5 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
228
tests/unit/proxy/guardrails/test_guardrail_helpers.py
Normal file
228
tests/unit/proxy/guardrails/test_guardrail_helpers.py
Normal file
|
|
@ -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 - {")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue