Merge pull request #41571 from BerriAI/litellm_policy_attachment_priority

feat(policy_engine): explicit priority for policy attachment execution order
This commit is contained in:
yucheng-berri 2026-09-17 17:27:55 -07:00 committed by GitHub
commit a9bea4f64d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 376 additions and 12 deletions

View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_PolicyAttachmentTable" ADD COLUMN IF NOT EXISTS "priority" INTEGER;

View file

@ -1379,6 +1379,7 @@ model LiteLLM_PolicyAttachmentTable {
keys String[] @default([]) // Key aliases or patterns
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

View file

@ -34019,6 +34019,20 @@
"title": "Policy Name",
"type": "string"
},
"priority": {
"anyOf": [
{
"maximum": 2147483647.0,
"minimum": -2147483648.0,
"type": "integer"
},
{
"type": "null"
}
],
"description": "Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
"title": "Priority"
},
"scope": {
"anyOf": [
{
@ -34132,6 +34146,18 @@
"title": "Policy Name",
"type": "string"
},
"priority": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
],
"description": "Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
"title": "Priority"
},
"scope": {
"anyOf": [
{
@ -36152,6 +36178,20 @@
"title": "Policy Name",
"type": "string"
},
"priority": {
"anyOf": [
{
"maximum": 2147483647.0,
"minimum": -2147483648.0,
"type": "integer"
},
{
"type": "null"
}
],
"description": "Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
"title": "Priority"
},
"scope": {
"anyOf": [
{

View file

@ -48,6 +48,13 @@ def _attachment_specificity(attachment: PolicyAttachment) -> tuple[int, int]:
return (max(dims, default=0), len(dims))
def _attachment_sort_key(attachment: PolicyAttachment) -> tuple[int, int, int, int]:
specificity: Final = _attachment_specificity(attachment)
if attachment.priority is not None:
return (0, attachment.priority, *specificity)
return (1, 0, *specificity)
class AttachmentRegistry:
"""
In-memory registry for storing and managing policy attachments.
@ -111,6 +118,7 @@ class AttachmentRegistry:
keys=attachment_data.get("keys"),
models=attachment_data.get("models"),
tags=attachment_data.get("tags"),
priority=attachment_data.get("priority"),
)
def get_attached_policies(self, context: PolicyMatchContext) -> list[str]:
@ -140,7 +148,7 @@ class AttachmentRegistry:
for attachment in self._attachments
if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
),
key=_attachment_specificity,
key=_attachment_sort_key,
)
broadest_attachment_by_policy: Final = MappingProxyType(
{attachment.policy: attachment for attachment in reversed(matching_attachments)}
@ -315,6 +323,7 @@ class AttachmentRegistry:
"keys": attachment_request.keys or [],
"models": attachment_request.models or [],
"tags": attachment_request.tags or [],
"priority": attachment_request.priority,
"created_at": datetime.now(timezone.utc),
"updated_at": datetime.now(timezone.utc),
"created_by": created_by,
@ -330,6 +339,7 @@ class AttachmentRegistry:
keys=attachment_request.keys,
models=attachment_request.models,
tags=attachment_request.tags,
priority=attachment_request.priority,
)
self.add_attachment(attachment)
@ -341,6 +351,7 @@ class AttachmentRegistry:
keys=created_attachment.keys or [],
models=created_attachment.models or [],
tags=created_attachment.tags or [],
priority=created_attachment.priority,
created_at=created_attachment.created_at,
updated_at=created_attachment.updated_at,
created_by=created_attachment.created_by,
@ -417,6 +428,7 @@ class AttachmentRegistry:
keys=attachment.keys or [],
models=attachment.models or [],
tags=attachment.tags or [],
priority=attachment.priority,
created_at=attachment.created_at,
updated_at=attachment.updated_at,
created_by=attachment.created_by,
@ -455,6 +467,7 @@ class AttachmentRegistry:
keys=a.keys or [],
models=a.models or [],
tags=a.tags or [],
priority=a.priority,
created_at=a.created_at,
updated_at=a.updated_at,
created_by=a.created_by,
@ -488,6 +501,7 @@ class AttachmentRegistry:
keys=attachment_response.keys if attachment_response.keys else None,
models=(attachment_response.models if attachment_response.models else None),
tags=attachment_response.tags if attachment_response.tags else None,
priority=attachment_response.priority,
)
for attachment_response in attachments
]

View file

@ -60,6 +60,7 @@ def _config_attachment_to_db_response(index: int, attachment: PolicyAttachment)
keys=attachment.keys or [],
models=attachment.models or [],
tags=attachment.tags or [],
priority=attachment.priority,
definition_location="config",
)

View file

@ -1379,6 +1379,7 @@ model LiteLLM_PolicyAttachmentTable {
keys String[] @default([]) // Key aliases or patterns
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

View file

@ -288,6 +288,12 @@ class PolicyAttachment(BaseModel):
default=None,
description="Tag patterns this attachment applies to. Supports wildcards (e.g., health-*).",
)
priority: int | None = Field(
default=None,
ge=-2147483648,
le=2147483647,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
model_config = ConfigDict(extra="forbid")

View file

@ -305,6 +305,12 @@ class PolicyAttachmentCreateRequest(BaseModel):
default=None,
description="Tag patterns this attachment applies to. Supports wildcards (e.g., health-*).",
)
priority: int | None = Field(
default=None,
ge=-2147483648,
le=2147483647,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
class PolicyAttachmentDBResponse(BaseModel):
@ -317,6 +323,10 @@ class PolicyAttachmentDBResponse(BaseModel):
keys: list[str] = Field(default_factory=list, description="Key patterns.")
models: list[str] = Field(default_factory=list, description="Model patterns.")
tags: list[str] = Field(default_factory=list, description="Tag patterns.")
priority: int | None = Field(
default=None,
description="Explicit execution order, lower runs first. Prioritised attachments run before those without one.",
)
created_at: datetime | None = Field(default=None, description="When the attachment was created.")
updated_at: datetime | None = Field(default=None, description="When the attachment was last updated.")
created_by: str | None = Field(default=None, description="Who created the attachment.")

View file

@ -1379,6 +1379,7 @@ model LiteLLM_PolicyAttachmentTable {
keys String[] @default([]) // Key aliases or patterns
models String[] @default([]) // Model names or patterns
tags String[] @default([]) // Tag patterns (e.g., ["healthcare", "prod-*"])
priority Int? // Explicit execution order
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt

View file

@ -158,6 +158,68 @@ class TestGetAttachedPolicies:
"model-policy",
]
def test_prioritized_attachments_run_before_unprioritized_attachments(self):
registry = AttachmentRegistry()
registry.load_attachments(
[
{"policy": "unprioritized-tag", "tags": ["prod"]},
{"policy": "prioritized-tag", "tags": ["prod"], "priority": 5},
{"policy": "prioritized-model", "models": ["gpt-4"], "priority": 0},
]
)
context = PolicyMatchContext(model="gpt-4", tags=["prod"])
assert registry.get_attached_policies(context) == [
"prioritized-model",
"prioritized-tag",
"unprioritized-tag",
]
def test_prioritized_attachments_order_by_priority_across_scope_tiers(self):
registry = AttachmentRegistry()
registry.load_attachments(
[
{"policy": "team-policy", "teams": ["team-a"], "priority": 2},
{"policy": "model-policy", "models": ["gpt-4"], "priority": 1},
]
)
context = PolicyMatchContext(team_alias="team-a", model="gpt-4")
assert registry.get_attached_policies(context) == ["model-policy", "team-policy"]
def test_equal_priority_attachments_fall_back_to_scope_tier_order(self):
registry = AttachmentRegistry()
registry.load_attachments(
[
{"policy": "model-policy", "models": ["gpt-4"], "priority": 1},
{"policy": "tag-policy", "tags": ["prod"], "priority": 1},
{"policy": "global-policy", "scope": "*", "priority": 1},
]
)
context = PolicyMatchContext(model="gpt-4", tags=["prod"])
assert registry.get_attached_policies(context) == ["global-policy", "tag-policy", "model-policy"]
def test_duplicate_policy_uses_highest_priority_attachment(self):
registry = AttachmentRegistry()
registry.load_attachments(
[
{"policy": "shared-policy", "scope": "*"},
{"policy": "global-policy", "scope": "*"},
{"policy": "shared-policy", "models": ["gpt-4"], "priority": 0},
]
)
context = PolicyMatchContext(model="gpt-4")
assert registry.get_attached_policies_with_reasons(context) == [
{"policy_name": "shared-policy", "matched_via": "model:gpt-4"},
{"policy_name": "global-policy", "matched_via": "scope:*"},
]
def test_combined_team_and_model_attachment_uses_model_specificity(self):
registry = AttachmentRegistry()
registry.load_attachments(
@ -474,8 +536,28 @@ class TestAttachmentRegistrySingleton:
registry2 = get_attachment_registry()
assert registry1 is registry2
def test_parse_attachment_reads_priority(self):
registry = AttachmentRegistry()
registry.load_attachments(
[
{"policy": "prioritized", "priority": 4},
{"policy": "unprioritized"},
]
)
def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scope=None, teams=None):
attachments = registry.get_all_attachments()
assert attachments[0].priority == 4
assert attachments[1].priority is None
def _make_db_attachment_row(
attachment_id: str = "att-1",
policy_name: str = "db-policy",
scope: str | None = None,
teams: list[str] | None = None,
priority: int | None = None,
) -> MagicMock:
row = MagicMock()
row.attachment_id = attachment_id
row.policy_name = policy_name
@ -484,6 +566,7 @@ def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scop
row.keys = []
row.models = []
row.tags = []
row.priority = priority
row.created_at = datetime.now(timezone.utc)
row.updated_at = datetime.now(timezone.utc)
row.created_by = None
@ -491,9 +574,11 @@ def _make_db_attachment_row(attachment_id="att-1", policy_name="db-policy", scop
return row
def _prisma_with_attachment_rows(rows):
def _prisma_with_attachment_rows(rows: list[MagicMock]) -> MagicMock:
prisma = MagicMock()
prisma.db.litellm_policyattachmenttable.find_many = AsyncMock(return_value=rows)
prisma.configure_mock(
**{"db.litellm_policyattachmenttable.find_many": AsyncMock(return_value=rows)}
)
return prisma
@ -535,6 +620,15 @@ class TestConfigAttachmentsPreservedAcrossDbSync:
assert len(registry.get_all_attachments()) == 1
@pytest.mark.asyncio
async def test_sync_round_trips_db_attachment_priority(self):
registry = AttachmentRegistry()
db_row = _make_db_attachment_row(priority=7)
await registry.sync_attachments_from_db(_prisma_with_attachment_rows([db_row]))
assert registry.get_all_attachments()[0].priority == 7
@pytest.mark.asyncio
async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
registry = AttachmentRegistry()

View file

@ -0,0 +1,15 @@
import pytest
from pydantic import ValidationError
from litellm.types.proxy.policy_engine.policy_types import PolicyAttachment
@pytest.mark.parametrize("priority", [-2147483648, 2147483647])
def test_policy_attachment_accepts_int32_priority(priority: int):
assert PolicyAttachment(policy="p", priority=priority).priority == priority
@pytest.mark.parametrize("priority", [-2147483649, 2147483648])
def test_policy_attachment_rejects_priority_outside_int32(priority: int):
with pytest.raises(ValidationError):
PolicyAttachment(policy="p", priority=priority)

View file

@ -3,8 +3,10 @@ Tests for pipeline field on policy CRUD types (resolver_types.py).
"""
import pytest
from pydantic import ValidationError
from litellm.types.proxy.policy_engine.resolver_types import (
PolicyAttachmentCreateRequest,
PolicyCreateRequest,
PolicyDBResponse,
PolicyUpdateRequest,
@ -100,3 +102,14 @@ def test_policy_create_request_roundtrip():
dumped = req.model_dump()
restored = PolicyCreateRequest(**dumped)
assert restored.pipeline == pipeline_data
@pytest.mark.parametrize("priority", [-2147483648, 2147483647])
def test_policy_attachment_create_request_accepts_int32_priority(priority: int):
assert PolicyAttachmentCreateRequest(policy_name="p", priority=priority).priority == priority
@pytest.mark.parametrize("priority", [-2147483649, 2147483648])
def test_policy_attachment_create_request_rejects_priority_outside_int32(priority: int):
with pytest.raises(ValidationError):
PolicyAttachmentCreateRequest(policy_name="p", priority=priority)

View file

@ -45,9 +45,26 @@ describe("AttachmentTable", () => {
expect(screen.getByText("Keys")).toBeInTheDocument();
expect(screen.getByText("Models")).toBeInTheDocument();
expect(screen.getByText("Tags")).toBeInTheDocument();
expect(screen.getByText("Priority")).toBeInTheDocument();
expect(screen.getByText("Created At")).toBeInTheDocument();
});
it("should show the priority and a dash for attachments without one", () => {
const attachments = [
makeAttachment({ attachment_id: "att-prio0001", policy_name: "prioritized", priority: 5 }),
makeAttachment({ attachment_id: "att-prio0002", policy_name: "unprioritized" }),
];
renderWithProviders(<AttachmentTable {...defaultProps} attachments={attachments} />);
const rows = screen.getAllByRole("row").slice(1);
const prioritizedRow = rows.find((row) => within(row).queryByText("prioritized"));
const unprioritizedRow = rows.find((row) => within(row).queryByText("unprioritized"));
expect(within(prioritizedRow!).getByText("5")).toBeInTheDocument();
expect(within(unprioritizedRow!).queryByText("5")).not.toBeInTheDocument();
expect(within(unprioritizedRow!).getAllByText("-")).toHaveLength(
within(prioritizedRow!).getAllByText("-").length + 1,
);
});
it("should show skeleton rows when isLoading is true", () => {
renderWithProviders(<AttachmentTable {...defaultProps} isLoading />);
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);

View file

@ -167,6 +167,20 @@ export const getAttachmentTableColumns = ({
enableSorting: false,
cell: ({ row }) => <ChipList values={row.original.tags ?? []} />,
},
{
id: "priority",
accessorFn: (row) => row.priority ?? Number.POSITIVE_INFINITY,
meta: { title: "Priority" },
header: ({ column }) => <DataTableSortHeader column={column} title="Priority" />,
size: 100,
enableSorting: true,
cell: ({ row }) =>
row.original.priority == null ? (
<span className="text-muted-foreground">-</span>
) : (
<span className="font-mono text-xs">{row.original.priority}</span>
),
},
{
id: "created_at",
accessorFn: (row) => row.created_at ?? "",

View file

@ -1,5 +1,5 @@
import React from "react";
import { screen, waitFor } from "@testing-library/react";
import { fireEvent, screen, waitFor } 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";
@ -180,6 +180,78 @@ describe("AddAttachmentForm", () => {
expect(screen.queryByText(TEAMS_ERROR)).not.toBeInTheDocument();
});
const selectPolicy = async (user: UserEvent, policyName: string) => {
await screen.findByText("Create Policy Attachment");
const input = screen.getByLabelText("Policies");
await user.click(input);
await user.type(input, `${policyName}{Enter}`);
};
const setPriority = (value: string) => {
fireEvent.change(screen.getByLabelText("Priority"), { target: { value } });
};
const submit = async (user: UserEvent) => {
await user.click(screen.getByRole("button", { name: /create attachment/i }));
};
it("sends the entered priority with the attachment", async () => {
const user = userEvent.setup();
const createAttachment = vi.fn().mockResolvedValue({});
renderWithProviders(<AddAttachmentForm {...defaultProps} createAttachment={createAttachment} />);
await selectPolicy(user, "policy-alpha");
setPriority("10");
await submit(user);
await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1));
expect(createAttachment).toHaveBeenCalledWith("test-token", {
policy_name: "policy-alpha",
scope: "*",
priority: 10,
});
});
it("sends a negative priority typed one keystroke at a time", async () => {
const user = userEvent.setup();
const createAttachment = vi.fn().mockResolvedValue({});
renderWithProviders(<AddAttachmentForm {...defaultProps} createAttachment={createAttachment} />);
await selectPolicy(user, "policy-alpha");
const priority = screen.getByLabelText("Priority");
await user.type(priority, "-5");
expect(priority).toHaveValue(-5);
await submit(user);
await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1));
expect(createAttachment).toHaveBeenCalledWith("test-token", {
policy_name: "policy-alpha",
scope: "*",
priority: -5,
});
});
it("omits priority from the attachment when the field is left blank", async () => {
const user = userEvent.setup();
const createAttachment = vi.fn().mockResolvedValue({});
renderWithProviders(<AddAttachmentForm {...defaultProps} createAttachment={createAttachment} />);
await selectPolicy(user, "policy-alpha");
await submit(user);
await waitFor(() => expect(createAttachment).toHaveBeenCalledTimes(1));
expect(createAttachment).toHaveBeenCalledWith("test-token", { policy_name: "policy-alpha", scope: "*" });
});
it.each([
["2147483648", /at most 2147483647/i],
["-2147483649", /at least -2147483648/i],
["1.5", /whole number/i],
])("blocks submit with a field error when priority is %s", async (value, error) => {
const user = userEvent.setup();
const createAttachment = vi.fn();
renderWithProviders(<AddAttachmentForm {...defaultProps} createAttachment={createAttachment} />);
await selectPolicy(user, "policy-alpha");
setPriority(value);
await submit(user);
expect(await screen.findByText(error)).toBeInTheDocument();
expect(createAttachment).not.toHaveBeenCalled();
});
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"));

View file

@ -8,6 +8,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { FieldGroup, FieldLabel, FieldTitle } from "@/components/ui/field";
import { FormField } from "@/components/shared/form/FormField";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
import { Separator } from "@/components/ui/separator";
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
@ -36,6 +37,7 @@ interface AttachmentFormValues {
keys: string[];
models: string[];
tags: string[];
priority: number | null;
}
const EMPTY_VALUES: AttachmentFormValues = {
@ -44,14 +46,24 @@ const EMPTY_VALUES: AttachmentFormValues = {
keys: [],
models: [],
tags: [],
priority: null,
};
const INT32_MIN = -2147483648;
const INT32_MAX = 2147483647;
const attachmentShape = {
policy_names: z.array(z.string()).min(1, "Please select at least one policy"),
teams: z.array(z.string()),
keys: z.array(z.string()),
models: z.array(z.string()),
tags: z.array(z.string()),
priority: z
.number({ error: "Priority must be a whole number" })
.int("Priority must be a whole number")
.min(INT32_MIN, `Priority must be at least ${INT32_MIN}`)
.max(INT32_MAX, `Priority must be at most ${INT32_MAX}`)
.nullable(),
};
const buildAttachmentSchema = (scopeType: ScopeType, teamsLoaded: boolean, availableTeams: string[]) =>
@ -419,6 +431,28 @@ const AddAttachmentForm: React.FC<AddAttachmentFormProps> = ({
</FormField>
</>
)}
<FormField
control={form.control}
name="priority"
label={labelWithHint(
"Priority",
"Lower numbers run first. Attachments with a priority run before attachments without one.",
)}
description="Optional. Leave blank to keep the default order: global, then teams, keys, tags, models."
>
{({ ref, value, onChange, ...field }) => (
<Input
{...field}
ref={ref}
type="number"
step={1}
value={value ?? ""}
placeholder="e.g. 10"
onChange={(event) => onChange(event.target.value === "" ? null : event.target.valueAsNumber)}
/>
)}
</FormField>
</FieldGroup>
{impactResult && <ImpactPreviewAlert impactResult={impactResult} />}

View file

@ -79,4 +79,18 @@ describe("buildAttachmentData", () => {
expect(result.tags).toBeUndefined();
});
});
describe("priority", () => {
it.each(["global", "specific"] as const)("should include priority for a %s scope", (scopeType) => {
expect(buildAttachmentData({ policy_name: "p", priority: 0 }, scopeType).priority).toBe(0);
});
it("should include a negative priority", () => {
expect(buildAttachmentData({ policy_name: "p", priority: -5 }, "specific").priority).toBe(-5);
});
it.each([undefined, null])("should omit priority when it is %s", (priority) => {
expect(buildAttachmentData({ policy_name: "p", priority }, "specific")).not.toHaveProperty("priority");
});
});
});

View file

@ -1,13 +1,16 @@
import { PolicyAttachmentCreateRequest } from "@/components/policies/types";
/**
* Builds a PolicyAttachmentCreateRequest from form values.
*
* @param formValues - The raw form field values (from form.getFieldsValue)
* @param scopeType - Whether the attachment is "global" or "specific"
*/
export interface AttachmentFormInput {
policy_name: string;
teams?: string[];
keys?: string[];
models?: string[];
tags?: string[];
priority?: number | null;
}
export function buildAttachmentData(
formValues: Record<string, any>,
formValues: AttachmentFormInput,
scopeType: "global" | "specific",
): PolicyAttachmentCreateRequest {
const data: PolicyAttachmentCreateRequest = {
@ -21,5 +24,6 @@ export function buildAttachmentData(
if (formValues.models && formValues.models.length > 0) data.models = formValues.models;
if (formValues.tags && formValues.tags.length > 0) data.tags = formValues.tags;
}
if (typeof formValues.priority === "number") data.priority = formValues.priority;
return data;
}

View file

@ -44,6 +44,7 @@ export interface PolicyAttachment {
keys: string[];
models: string[];
tags: string[];
priority?: number | null;
created_at?: string;
updated_at?: string;
created_by?: string;
@ -78,6 +79,7 @@ export interface PolicyAttachmentCreateRequest {
keys?: string[];
models?: string[];
tags?: string[];
priority?: number;
}
export interface PolicyListResponse {

View file

@ -34566,6 +34566,11 @@ export interface components {
* @description Name of the policy to attach.
*/
policy_name: string;
/**
* Priority
* @description Explicit execution order, lower runs first. Prioritised attachments run before those without one.
*/
priority?: number | null;
/**
* Scope
* @description Use '*' for global scope (applies to all requests).
@ -34624,6 +34629,11 @@ export interface components {
* @description Name of the attached policy.
*/
policy_name: string;
/**
* Priority
* @description Explicit execution order, lower runs first. Prioritised attachments run before those without one.
*/
priority?: number | null;
/**
* Scope
* @description Scope of the attachment.