diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 4a94fc3d43c3..a879f6b6f7ee 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -695,7 +695,8 @@ async def create_policy_attachment( } ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.policy_engine.policy_validator import PolicyValidator + from litellm.proxy.proxy_server import llm_router, prisma_client if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") @@ -710,6 +711,19 @@ async def create_policy_attachment( detail=f"Policy '{request.policy_name}' not found. Create the policy first.", ) + # Reject concrete team/key/model scope entries that don't resolve to a real + # entity. Wildcard patterns are allowed through (they may match zero today). + scope_errors = await PolicyValidator( + prisma_client=prisma_client, llm_router=llm_router + ).find_invalid_scope_entries( + policy_name=request.policy_name, + teams=request.teams, + keys=request.keys, + models=request.models, + ) + if scope_errors: + raise HTTPException(status_code=400, detail=" | ".join(e.message for e in scope_errors)) + created_by = user_api_key_dict.user_id result = await get_attachment_registry().add_attachment_to_db( attachment_request=request, diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py index 626bbbc1ce50..4db5bc0435e5 100644 --- a/litellm/proxy/policy_engine/policy_validator.py +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -9,9 +9,11 @@ - Inheritance chains are valid (no cycles, parents exist) """ +import asyncio from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.route_checks import RouteChecks from litellm.repositories.team_repository import TeamRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, @@ -152,6 +154,71 @@ def check_model_exists(self, model: str) -> bool: verbose_proxy_logger.warning(f"Could not check model '{model}': {str(e)}") return True # Assume valid on error + @staticmethod + def _scope_error( + policy_name: str, + error_type: PolicyValidationErrorType, + field: str, + value: str, + label: str, + ) -> PolicyValidationError: + return PolicyValidationError( + policy_name=policy_name, + error_type=error_type, + message=( + f"{label.capitalize()} '{value}' does not exist. Reference an existing " + f"{label} or use a wildcard pattern (e.g. '{value}*') to match by prefix." + ), + field=field, + value=value, + ) + + async def find_invalid_scope_entries( + self, + policy_name: str, + teams: list[str] | None = None, + keys: list[str] | None = None, + models: list[str] | None = None, + ) -> list[PolicyValidationError]: + """ + Validate the concrete scope entries of a policy attachment. + + Returns an error for every non-wildcard entry that does not resolve to an + existing team alias, key alias, or model. Wildcard patterns are always + accepted: a pattern like "healthcare-*" may match zero entities today and + match ones created later, so it cannot be validated by existence. Tags are + intentionally not checked - they are free-form labels with no registry to + validate against. + """ + # A concrete entry is one the request-time matcher compares by exact equality; + # only a trailing "*" is a wildcard (RouteChecks._is_wildcard_pattern), and those + # are left unvalidated since they may match zero entities today and more later. + is_pattern = RouteChecks._is_wildcard_pattern + concrete_teams = [t for t in (teams or []) if not is_pattern(pattern=t)] + concrete_keys = [k for k in (keys or []) if not is_pattern(pattern=k)] + concrete_models = [m for m in (models or []) if not is_pattern(pattern=m)] + + team_exists = await asyncio.gather(*(self.check_team_alias_exists(t) for t in concrete_teams)) + key_exists = await asyncio.gather(*(self.check_key_alias_exists(k) for k in concrete_keys)) + + return [ + *( + self._scope_error(policy_name, PolicyValidationErrorType.INVALID_TEAM, "teams", team, "team") + for team, exists in zip(concrete_teams, team_exists) + if not exists + ), + *( + self._scope_error(policy_name, PolicyValidationErrorType.INVALID_KEY, "keys", key, "key") + for key, exists in zip(concrete_keys, key_exists) + if not exists + ), + *( + self._scope_error(policy_name, PolicyValidationErrorType.INVALID_MODEL, "models", model, "model") + for model in concrete_models + if not self.check_model_exists(model) + ), + ] + def _validate_inheritance_chain( self, policy_name: str, diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_validator.py b/tests/test_litellm/proxy/policy_engine/test_policy_validator.py index de2dde586980..a56695ee7f5c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_validator.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_validator.py @@ -6,6 +6,7 @@ - Guardrail names exist in registry """ +from typing import Optional, Set from unittest.mock import MagicMock, patch import pytest @@ -18,6 +19,37 @@ ) +class _FakeTable: + """Minimal stand-in for a Prisma table whose find_first matches on one field.""" + + def __init__(self, existing: Set[str], match_field: str): + self._existing = existing + self._match_field = match_field + + async def find_first(self, where: dict) -> Optional[object]: + return object() if where.get(self._match_field) in self._existing else None + + +class _FakeDB: + def __init__(self, teams: Set[str], keys: Set[str]): + self.litellm_teamtable = _FakeTable(teams, "team_alias") + self.litellm_verificationtoken = _FakeTable(keys, "key_alias") + + +class _FakePrisma: + """Injected prisma client so PolicyValidator can be unit-tested without a DB.""" + + def __init__(self, teams: Set[str] = frozenset(), keys: Set[str] = frozenset()): + self.db = _FakeDB(teams, keys) + + +class _FakeRouter: + """Injected router exposing only what check_model_exists reads.""" + + def __init__(self, model_names: Set[str]): + self.model_names = list(model_names) + + class TestPolicyValidator: """Test policy validation logic.""" @@ -91,3 +123,72 @@ async def test_validate_valid_policy(self): assert result.valid is True assert len(result.errors) == 0 + + +class TestAttachmentScopeValidation: + """Regression tests for LIT-4199: attachments must not accept non-existent teams/keys/models.""" + + @pytest.mark.asyncio + async def test_nonexistent_team_is_flagged(self): + validator = PolicyValidator(prisma_client=_FakePrisma(teams={"real-team"})) + errors = await validator.find_invalid_scope_entries( + policy_name="p", teams=["real-team", "ghost-team"] + ) + assert [(e.field, e.value) for e in errors] == [("teams", "ghost-team")] + assert errors[0].error_type == PolicyValidationErrorType.INVALID_TEAM + + @pytest.mark.asyncio + async def test_existing_team_passes(self): + validator = PolicyValidator(prisma_client=_FakePrisma(teams={"payments"})) + errors = await validator.find_invalid_scope_entries(policy_name="p", teams=["payments"]) + assert errors == [] + + @pytest.mark.asyncio + async def test_trailing_star_wildcard_is_allowed_even_when_it_matches_nothing(self): + validator = PolicyValidator(prisma_client=_FakePrisma(teams=set())) + errors = await validator.find_invalid_scope_entries( + policy_name="p", teams=["healthcare-*", "brand-new-*"] + ) + assert errors == [] + + @pytest.mark.asyncio + async def test_only_trailing_star_counts_as_a_wildcard(self): + # Request-time matching treats only a trailing "*" as a wildcard; "?" and a + # non-trailing "*" are compared literally, so they are validated as concrete + # aliases (and here resolve to nothing -> flagged). + validator = PolicyValidator(prisma_client=_FakePrisma(teams=set())) + errors = await validator.find_invalid_scope_entries( + policy_name="p", teams=["ops-?", "heal*care"] + ) + assert {e.value for e in errors} == {"ops-?", "heal*care"} + + @pytest.mark.asyncio + async def test_keys_and_models_are_validated_too(self): + validator = PolicyValidator( + prisma_client=_FakePrisma(keys={"prod-key"}), + llm_router=_FakeRouter(model_names={"gpt-4o"}), + ) + errors = await validator.find_invalid_scope_entries( + policy_name="p", + keys=["prod-key", "ghost-key"], + models=["gpt-4o", "ghost-model", "bedrock/*"], + ) + flagged = {(e.field, e.value) for e in errors} + assert ("keys", "ghost-key") in flagged + assert ("models", "ghost-model") in flagged + assert ("keys", "prod-key") not in flagged + assert ("models", "gpt-4o") not in flagged + assert ("models", "bedrock/*") not in flagged # wildcard model allowed through + + @pytest.mark.asyncio + async def test_no_scope_entries_returns_no_errors(self): + validator = PolicyValidator(prisma_client=_FakePrisma()) + errors = await validator.find_invalid_scope_entries(policy_name="p") + assert errors == [] + + @pytest.mark.asyncio + async def test_without_db_connection_assumes_valid(self): + # Fail-open: with no DB we cannot verify existence, so nothing is blocked. + validator = PolicyValidator(prisma_client=None) + errors = await validator.find_invalid_scope_entries(policy_name="p", teams=["anything"]) + assert errors == [] diff --git a/ui/litellm-dashboard/src/components/policies/add_attachment_form.test.tsx b/ui/litellm-dashboard/src/components/policies/add_attachment_form.test.tsx index fa02c85841dd..adbbb1b6997b 100644 --- a/ui/litellm-dashboard/src/components/policies/add_attachment_form.test.tsx +++ b/ui/litellm-dashboard/src/components/policies/add_attachment_form.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "../../../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -9,6 +9,10 @@ import { Policy } from "./types"; vi.mock("../networking"); +vi.mock("../molecules/notifications_manager", () => ({ + default: { success: vi.fn(), fromBackend: vi.fn(), error: vi.fn(), info: vi.fn() }, +})); + vi.mock("./impact_preview_alert", () => ({ default: ({ impactResult }: { impactResult: any }) => React.createElement("div", { "data-testid": "impact-preview" }, `${impactResult.affected_keys_count} keys`), @@ -37,6 +41,9 @@ const defaultProps = { createAttachment: vi.fn(), }; +const teamListResult = (aliases: string[]) => + aliases.map((team_alias) => ({ team_alias })) as unknown as Awaited>; + describe("AddAttachmentForm", () => { beforeEach(() => { vi.clearAllMocks(); @@ -110,4 +117,71 @@ describe("AddAttachmentForm", () => { renderWithProviders(); expect(await screen.findByRole("button", { name: /create attachment/i })).toBeInTheDocument(); }); + + type UserEvent = ReturnType; + + const TEAMS_ERROR = /these teams don't exist/i; + + const openSpecificScope = async (user: UserEvent) => { + await screen.findByText("Create Policy Attachment"); + await waitFor(() => expect(networking.teamListCall).toHaveBeenCalled()); + await user.click(screen.getByRole("radio", { name: /specific/i })); + }; + + const enterTeam = async (user: UserEvent, value: string) => { + const item = screen.getByText("Teams").closest(".ant-form-item") as HTMLElement; + const input = within(item).getByRole("combobox"); + await user.click(input); + await user.type(input, `${value}{Enter}`); + }; + + // Submits and waits for the validation cycle to settle. No policy is selected, so + // the "select at least one policy" required error always appears - we use it as a + // synchronization point, then assert whether the teams validator also complained. + const submitAndSettle = async (user: UserEvent) => { + await user.click(screen.getByRole("button", { name: /create attachment/i })); + await screen.findByText(/select at least one policy/i); + }; + + it("blocks submit with a field error when a concrete team that does not exist is entered", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamListCall).mockResolvedValue(teamListResult(["real-team"])); + const createAttachment = vi.fn(); + renderWithProviders(); + await openSpecificScope(user); + await enterTeam(user, "ghost-team"); + await submitAndSettle(user); + expect(screen.getByText(TEAMS_ERROR)).toBeInTheDocument(); + expect(createAttachment).not.toHaveBeenCalled(); + }); + + it("does not flag a team that exists", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamListCall).mockResolvedValue(teamListResult(["real-team"])); + renderWithProviders(); + await openSpecificScope(user); + await enterTeam(user, "real-team"); + await submitAndSettle(user); + expect(screen.queryByText(TEAMS_ERROR)).not.toBeInTheDocument(); + }); + + it("does not flag a wildcard pattern even when it matches no existing team", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamListCall).mockResolvedValue(teamListResult([])); + renderWithProviders(); + await openSpecificScope(user); + await enterTeam(user, "healthcare-*"); + await submitAndSettle(user); + expect(screen.queryByText(TEAMS_ERROR)).not.toBeInTheDocument(); + }); + + it("defers to the backend (does not flag) when the team list failed to load", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamListCall).mockRejectedValue(new Error("boom")); + renderWithProviders(); + await openSpecificScope(user); + await enterTeam(user, "ghost-team"); + await submitAndSettle(user); + expect(screen.queryByText(TEAMS_ERROR)).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/policies/add_attachment_form.tsx b/ui/litellm-dashboard/src/components/policies/add_attachment_form.tsx index e19a14a2b825..dd4d07b3f359 100644 --- a/ui/litellm-dashboard/src/components/policies/add_attachment_form.tsx +++ b/ui/litellm-dashboard/src/components/policies/add_attachment_form.tsx @@ -6,6 +6,7 @@ import { teamListCall, keyListCall, modelAvailableCall, estimateAttachmentImpact import NotificationsManager from "../molecules/notifications_manager"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { buildAttachmentData } from "./build_attachment_data"; +import { getInvalidTeamEntries } from "./scope_validation"; import ImpactPreviewAlert from "./impact_preview_alert"; const { Text } = Typography; @@ -31,6 +32,7 @@ const AddAttachmentForm: React.FC = ({ const [isSubmitting, setIsSubmitting] = useState(false); const [scopeType, setScopeType] = useState<"global" | "specific">("global"); const [availableTeams, setAvailableTeams] = useState([]); + const [teamsLoaded, setTeamsLoaded] = useState(false); const [availableKeys, setAvailableKeys] = useState([]); const [availableModels, setAvailableModels] = useState([]); const [isLoadingTeams, setIsLoadingTeams] = useState(false); @@ -52,11 +54,13 @@ const AddAttachmentForm: React.FC = ({ // Load teams — teamListCall returns a plain array of team objects setIsLoadingTeams(true); + setTeamsLoaded(false); try { const teamsResponse = await teamListCall(accessToken, null, userId); const teamsArray = Array.isArray(teamsResponse) ? teamsResponse : teamsResponse?.data || []; const teamAliases = teamsArray.map((t: any) => t.team_alias).filter(Boolean); setAvailableTeams(teamAliases); + setTeamsLoaded(true); } catch (error) { console.error("Failed to load teams:", error); } finally { @@ -226,6 +230,20 @@ const AddAttachmentForm: React.FC = ({ name="teams" label="Teams" tooltip="Select team aliases or enter custom patterns. Supports wildcards (e.g., healthcare-*)" + rules={[ + { + validator: async (_rule, value?: string[]) => { + if (!teamsLoaded) return; + const invalid = getInvalidTeamEntries(value ?? [], availableTeams); + if (invalid.length > 0) { + throw new Error( + `These teams don't exist: ${invalid.join(", ")}. ` + + `Choose an existing team, or use a wildcard like "team-*" to match by prefix.`, + ); + } + }, + }, + ]} >