From 540caa6574eb4040f2c114236aea26de1f3afe88 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 15 Aug 2026 15:30:28 -0700 Subject: [PATCH 1/2] feat(ui): direction picker and reverse-mode display for shadow evals (#36994) * feat(ui): direction picker and reverse-mode display for shadow evals * fix(ui): include configured model groups in the shadow eval baseline picker --- .../_components/ShadowEvalSection.test.tsx | 55 +++++ .../_components/ShadowEvalSection.tsx | 194 ++++++++++++++---- .../hooks/models/useModels.test.ts | 27 +++ .../app/(dashboard)/hooks/models/useModels.ts | 21 ++ 4 files changed, 253 insertions(+), 44 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx index d4d26650086..e342ee33f25 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.test.tsx @@ -46,6 +46,7 @@ vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ { model_name: "gpt-auto", litellm_params: { model: "auto_router/gpt-auto" } }, ], })), + usePlainModelGroups: vi.fn(() => new Set(["prod-claude"])), })); vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ @@ -72,6 +73,8 @@ const job = (overrides: Partial = {}): ShadowEvalJob => ({ job_id: "job-1", status: "running", router_name: "claude-auto", + direction: "forward", + baseline_model: null, judge_model: "anthropic/claude-sonnet-5", shadow_percentage: 10, max_turns: 200, @@ -367,6 +370,58 @@ describe("ShadowEvalSection", () => { expect(start.mutate).toHaveBeenCalledWith(expectedBody); }); + it("requires a baseline model in reverse mode and submits it, while forward mode never shows the picker", async () => { + const user = userEvent.setup(); + const { start } = mockHooks({}); + render(); + + expect(screen.queryByPlaceholderText("Select a baseline model")).not.toBeInTheDocument(); + + await user.click(screen.getByText("Adoption check: key's traffic vs the router")); + await user.click(await screen.findByText("Regression check: router's picks vs a baseline")); + await user.click(screen.getByPlaceholderText("Search keys by alias")); + await user.click(await screen.findByText("prod-alpha")); + await user.click(screen.getByPlaceholderText("Select an auto-router")); + await user.click(await screen.findByText("gpt-auto")); + await user.click(screen.getByPlaceholderText("Select a judge model")); + await user.click(await screen.findByRole("option", { name: /anthropic\/claude-sonnet-5/ })); + + expect(screen.getByText("Start shadow eval")).toBeDisabled(); + + await user.click(screen.getByPlaceholderText("Select a baseline model")); + expect(await screen.findByRole("option", { name: /openai\/gpt-4o/ })).toBeInTheDocument(); + await user.click(screen.getByRole("option", { name: /prod-claude/ })); + await user.click(screen.getByText("Start shadow eval")); + + const expectedBody = { + api_key_id: "hash-alpha", + router_name: "gpt-auto", + direction: "reverse", + baseline_model: "prod-claude", + shadow_percentage: 10, + duration_days: 7, + max_turns: 200, + judge_model: "anthropic/claude-sonnet-5", + }; + expect(start.mutate).toHaveBeenCalledWith(expectedBody); + }); + + it("flips the arm labels and headline for a reverse job's results", () => { + const j = job({ direction: "reverse", baseline_model: "openai/gpt-4o" }); + mockHooks({ jobs: [j], detailsById: { "job-1": j } }); + render(); + + expect(screen.getByText(/on 10% of its traffic/)).toBeInTheDocument(); + expect(screen.getByText("Router matched or beat the baseline")).toBeInTheDocument(); + expect(screen.getByText("52.0%")).toBeInTheDocument(); + expect(screen.getByText(/Router won 30.0%/)).toBeInTheDocument(); + expect(screen.getByText(/Baseline won 48.0%/)).toBeInTheDocument(); + expect(screen.getAllByText("Baseline wins")).toHaveLength(2); + expect(screen.getByText("Router pick")).toBeInTheDocument(); + expect(screen.queryByText(/Current model/)).not.toBeInTheDocument(); + expect(screen.queryByText("Compared against")).not.toBeInTheDocument(); + }); + it("keeps an older job's verdicts reachable through the previous evaluations list", async () => { const user = userEvent.setup(); const emptyOverrides: Partial = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx index 711fc1af539..df2989d3990 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/ShadowEvalSection.tsx @@ -5,7 +5,7 @@ import React, { useMemo, useState } from "react"; import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; -import { useAutoRouters } from "@/app/(dashboard)/hooks/models/useModels"; +import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels"; import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect"; import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect"; import { Badge } from "@/components/ui/badge"; @@ -31,6 +31,37 @@ const pct = (value: number): string => `${value.toFixed(1)}%`; const MIN_TURNS_FOR_CONFIDENCE = 30; +type ShadowEvalDirection = ShadowEvalJob["direction"]; + +const otherArmLabel = (direction: ShadowEvalDirection): string => + direction === "reverse" ? "Baseline" : "Current model"; + +const routerWinRate = (direction: ShadowEvalDirection, slice: ShadowEvalSlice): number => + direction === "reverse" ? slice.real_win_rate_pct : slice.shadow_win_rate_pct; + +const otherArmWinRate = (direction: ShadowEvalDirection, slice: ShadowEvalSlice): number => + direction === "reverse" ? slice.shadow_win_rate_pct : slice.real_win_rate_pct; + +const routerMatchedOrBeatPct = ( + direction: ShadowEvalDirection, + results: NonNullable, +): number => + direction === "reverse" + ? 100 - results.overall_shadow_win_rate_pct + : results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct; + +const jobHeadline = (job: ShadowEvalJob): React.ReactNode => + job.direction === "reverse" ? ( + <> + Comparing {job.router_name} to{" "} + {job.baseline_model} on {job.shadow_percentage}% of its traffic + + ) : ( + <> + Shadowing {job.shadow_percentage}% via {job.router_name} + + ); + const isActive = (job: ShadowEvalJob): boolean => job.status === "running"; const endsIn = (endsAt: string | null | undefined): string | null => { @@ -54,16 +85,22 @@ const StatusBadge: React.FC<{ status: string }> = ({ status }) => ( ); -const SliceTable: React.FC<{ groupHeader: string; slices: readonly ShadowEvalSlice[] }> = ({ groupHeader, slices }) => ( +const SliceTable: React.FC<{ + groupHeader: string; + direction: ShadowEvalDirection; + slices: readonly ShadowEvalSlice[]; +}> = ({ groupHeader, direction, slices }) => ( {groupHeader} - {["Judged turns", "Router wins", "Current model wins", "Ties", "Judge confidence"].map((label) => ( - - {label} - - ))} + {["Judged turns", "Router wins", `${otherArmLabel(direction)} wins`, "Ties", "Judge confidence"].map( + (label) => ( + + {label} + + ), + )} @@ -77,9 +114,9 @@ const SliceTable: React.FC<{ groupHeader: string; slices: readonly ShadowEvalSli {slice.turn_count.toLocaleString()} - {pct(slice.shadow_win_rate_pct)} + {pct(routerWinRate(direction, slice))} - {pct(slice.real_win_rate_pct)} + {pct(otherArmWinRate(direction, slice))} {pct(slice.tie_rate_pct)} {slice.avg_judge_confidence.toFixed(2)} @@ -88,13 +125,23 @@ const SliceTable: React.FC<{ groupHeader: string; slices: readonly ShadowEvalSli
); -const VerdictBar: React.FC<{ results: NonNullable }> = ({ results }) => { - const routerWins = results.overall_shadow_win_rate_pct; +const VerdictBar: React.FC<{ direction: ShadowEvalDirection; results: NonNullable }> = ({ + direction, + results, +}) => { const ties = results.overall_tie_rate_pct; + const routerWins = + direction === "reverse" + ? Math.max(0, 100 - results.overall_shadow_win_rate_pct - ties) + : results.overall_shadow_win_rate_pct; const segments = [ { label: "Router won", value: routerWins, fill: "bg-emerald-500" }, { label: "Tie", value: ties, fill: "bg-emerald-200" }, - { label: "Current model won", value: Math.max(0, 100 - routerWins - ties), fill: "bg-muted-foreground/30" }, + { + label: `${otherArmLabel(direction)} won`, + value: Math.max(0, 100 - routerWins - ties), + fill: "bg-muted-foreground/30", + }, ]; return (
@@ -133,20 +180,22 @@ const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({ <>

- Router matched or beat your current model -

-

- {pct(results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct)} + Router matched or beat {job.direction === "reverse" ? "the baseline" : "your current model"}

+

{pct(routerMatchedOrBeatPct(job.direction, results))}

of {(job.judged_count ?? 0).toLocaleString()} judged responses

- + {results.by_current_model.length > 0 && ( - + )} {results.by_tier.length > 0 && (
0 ? "border-t" : ""}> - +
)} @@ -168,9 +217,7 @@ const JobResults: React.FC<{
-

- Shadowing {job.shadow_percentage}% via {job.router_name} -

+

{jobHeadline(job)}

{(job.judged_count ?? 0).toLocaleString()} of {job.max_turns.toLocaleString()} turns judged ·{" "} {(job.error_count ?? 0).toLocaleString()} errored · {usd(job.judge_spend ?? 0)} judge spend @@ -201,25 +248,55 @@ interface CostMapEntry { mode?: string; } -const useJudgeModelOptions = (): SearchSelectOption[] => { +const useChatModelNames = (): string[] => { const { data: costMap } = useModelCostMap(); + return useMemo(() => { + if (!costMap) return []; + const chatModels = Object.entries(costMap as Record) + .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) + .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); + return [...new Set(chatModels)].toSorted((a, b) => a.localeCompare(b)); + }, [costMap]); +}; + +const useJudgeModelOptions = (): SearchSelectOption[] => { + const chatModels = useChatModelNames(); return useMemo(() => { const pinned: SearchSelectOption[] = RECOMMENDED_JUDGE_MODELS.map((model) => ({ label: model, value: model, sublabel: "Recommended", })); - if (!costMap) return pinned; const pinnedNames = new Set(RECOMMENDED_JUDGE_MODELS); - const chatModels = Object.entries(costMap as Record) - .filter(([, value]) => value?.mode === "chat" && value?.litellm_provider) - .map(([key, value]) => (key.startsWith(`${value.litellm_provider}/`) ? key : `${value.litellm_provider}/${key}`)); - const rest = [...new Set(chatModels)] - .filter((model) => !pinnedNames.has(model)) - .toSorted((a, b) => a.localeCompare(b)) - .map((model) => ({ label: model, value: model })); + const rest = chatModels.filter((model) => !pinnedNames.has(model)).map((model) => ({ label: model, value: model })); return [...pinned, ...rest]; - }, [costMap]); + }, [chatModels]); +}; + +const useBaselineModelOptions = (): SearchSelectOption[] => { + const configuredGroups = usePlainModelGroups(); + const chatModels = useChatModelNames(); + return useMemo(() => { + const configured = [...configuredGroups] + .toSorted((a, b) => a.localeCompare(b)) + .map((model) => ({ label: model, value: model, sublabel: "Configured on this gateway" })); + const rest = chatModels + .filter((model) => !configuredGroups.has(model)) + .map((model) => ({ label: model, value: model })); + return [...configured, ...rest]; + }, [configuredGroups, chatModels]); +}; + +const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[] = [ + { value: "forward", label: "Adoption check: key's traffic vs the router" }, + { value: "reverse", label: "Regression check: router's picks vs a baseline" }, +] as const; + +const START_FORM_DESCRIPTION: Record = { + forward: + "Duplicates a sampled slice of the key's traffic through the auto-router and has an LLM judge compare both answers blind. The router's answers are never served to users; judge calls bill to the shadowed key.", + reverse: + "Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. The baseline's answers are never served to users; judge calls bill to the shadowed key.", }; const DURATION_OPTIONS = [ @@ -282,12 +359,15 @@ const StartForm: React.FC = () => { const { accessToken } = useAuthorized(); const [apiKeyId, setApiKeyId] = useState(""); const [routerName, setRouterName] = useState(""); + const [direction, setDirection] = useState("forward"); + const [baselineModel, setBaselineModel] = useState(""); const [percentage, setPercentage] = useState("10"); const [durationDays, setDurationDays] = useState("7"); const [judgeModel, setJudgeModel] = useState(""); const [maxTurns, setMaxTurns] = useState("200"); const { data: autoRouters } = useAutoRouters(); const judgeModelOptions = useJudgeModelOptions(); + const baselineModelOptions = useBaselineModelOptions(); const start = useStartShadowEval(); const routerOptions = useMemo(() => { @@ -301,14 +381,17 @@ const StartForm: React.FC = () => { const percentageValid = parsedPct >= 0.1 && parsedPct <= 100; const parsedMaxTurns = Number.parseInt(maxTurns, 10); const maxTurnsValid = parsedMaxTurns >= 1 && parsedMaxTurns <= 2000; - const filled = [apiKeyId, routerName, judgeModel].every((field) => field !== ""); + const filled = + [apiKeyId, routerName, judgeModel].every((field) => field !== "") && + (direction === "forward" || baselineModel !== ""); const boundsValid = percentageValid && maxTurnsValid; const valid = Boolean(accessToken) && filled && boundsValid; const handleStart = () => { const startBody = { api_key_id: apiKeyId, router_name: routerName, - direction: "forward" as const, + direction, + ...(direction === "reverse" ? { baseline_model: baselineModel } : {}), shadow_percentage: parsedPct, duration_days: Number.parseInt(durationDays, 10), max_turns: parsedMaxTurns, @@ -321,13 +404,27 @@ const StartForm: React.FC = () => { Start a shadow eval -

- Duplicates a sampled slice of the key's traffic through the auto-router and has an LLM judge compare both - answers blind. The router's answers are never served to users; judge calls bill to the shadowed key. -

+

{START_FORM_DESCRIPTION[direction]}

+ + + @@ -390,6 +487,17 @@ const StartForm: React.FC = () => {

Enter a value from 1 to 2000

)} + {direction === "reverse" && ( + + + + )} { const previousSummary = (job: ShadowEvalJob): string => { const results = job.results; - if (results) return pct(results.overall_shadow_win_rate_pct + results.overall_tie_rate_pct); + if (results) return pct(routerMatchedOrBeatPct(job.direction, results)); return job.judged_count === 0 ? "no verdicts" : "view results"; }; @@ -429,9 +537,7 @@ const PreviousJob: React.FC<{ job: ShadowEvalJob }> = ({ job }) => {
-

- {shown.shadow_percentage}% via {shown.router_name} -

+

{jobHeadline(shown)}

{shown.judged_count != null && `${shown.judged_count.toLocaleString()} judged · ${(shown.error_count ?? 0).toLocaleString()} errored · ${usd(shown.judge_spend ?? 0)} judge spend · `} @@ -507,8 +613,8 @@ const ShadowEvalSection: React.FC = () => {

Shadow eval

- Would the auto-router have answered as well as the models you use today? Find out on your real traffic, before - switching anything. + Blind-judge the auto-router on your real traffic: against the models a key uses today before switching, or + against a fixed baseline after it has switched.

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index f04c4b7bfcd..411e8402e11 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { isAutoRouterDeployment, selectAutoRouterModelGroups, + selectPlainModelGroups, useAllProxyModels, useAutoRouterModelGroups, useAutoRouters, @@ -977,6 +978,32 @@ describe("selectAutoRouterModelGroups", () => { }); }); +describe("selectPlainModelGroups", () => { + it("keeps only non-auto-router model groups", () => { + const deployments: AutoRouterCandidateDeployment[] = [ + { model_name: "smart-router", litellm_params: { model: "auto_router/complexity_router" } }, + { model_name: "claude-haiku", litellm_params: { model: "anthropic/claude-haiku-4-5" } }, + { model_name: "claude-sonnet", litellm_params: { model: "anthropic/claude-sonnet-4-5" } }, + { model_name: "cheap-router", litellm_params: { model: "auto_router/adaptive_router" } }, + ]; + + expect(selectPlainModelGroups(deployments)).toEqual(new Set(["claude-haiku", "claude-sonnet"])); + }); + + it("drops a group name that also fronts an auto-router deployment", () => { + const deployments: AutoRouterCandidateDeployment[] = [ + { model_name: "shared-name", litellm_params: { model: "auto_router/complexity_router" } }, + { model_name: "shared-name", litellm_params: { model: "anthropic/claude-sonnet-4-5" } }, + ]; + + expect(selectPlainModelGroups(deployments)).toEqual(new Set()); + }); + + it("drops deployments that have no public model_name", () => { + expect(selectPlainModelGroups([{ model_name: "", litellm_params: { model: "openai/gpt-4o" } }])).toEqual(new Set()); + }); +}); + describe("useAutoRouterModelGroups", () => { let queryClient: QueryClient; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index 68044a45edc..a5fbc433ea3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -123,6 +123,16 @@ export const selectAutoRouterModelGroups = (deployments: AutoRouterCandidateDepl export const selectAutoRouterDeployments = (deployments: AutoRouterDeployment[]): AutoRouterDeployment[] => deployments.filter(isAutoRouterDeployment); +export const selectPlainModelGroups = (deployments: AutoRouterCandidateDeployment[]): ReadonlySet => { + const autoRouterGroups = selectAutoRouterModelGroups(deployments); + return new Set( + deployments + .map((deployment) => deployment.model_name) + .filter((modelName): modelName is string => Boolean(modelName)) + .filter((modelName) => !autoRouterGroups.has(modelName)), + ); +}; + export const fetchAllModelDeployments = async ( accessToken: string, userId: string, @@ -172,6 +182,17 @@ export const useAutoRouterModelGroups = (): ReadonlySet => { return data ?? NO_AUTO_ROUTERS; }; +export const usePlainModelGroups = (): ReadonlySet => { + const { accessToken, userId, userRole } = useAuthorized(); + const { data } = useQuery>({ + queryKey: autoRouterListKey(userId, userRole), + queryFn: async () => await fetchAllModelDeployments(accessToken!, userId!, userRole!), + enabled: Boolean(accessToken && userId && userRole), + select: selectPlainModelGroups, + }); + return data ?? NO_AUTO_ROUTERS; +}; + export const useAutoRouters = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ From d7d10be0639395da813d6b9c88933bfafbcb3f85 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 15 Aug 2026 15:31:23 -0700 Subject: [PATCH 2/2] fix(guardrails): return the full PANW AIRS scan response on blocked requests (#37036) * fix(guardrails): return the full PANW AIRS scan response on blocked requests The blocked-request error detail was assembled from a hardcoded allowlist, so audit fields like prompt_detection_details, prompt_masked_data, source, transaction_id and session_id never reached the client even though AIRS returned them. Resolves LIT-5638 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(guardrails): drop redundant comment in AIRS error detail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(panw_prisma_airs): withhold response_masked_data from the blocked-response error The full AIRS passthrough also reached the response-side block path, where response_masked_data carries the model's own generation. That branch is only reached when mask_response_content is False, so the operator had explicitly declined to deliver that text, and the error body handed it back anyway. Withhold response_masked_data from the client-visible detail. prompt_masked_data stays: it is the caller's own input and one of the fields the ticket asks for. Every other AIRS field, including prompt_detection_details, source, transaction_id and session_id, is unchanged. * fix(panw_prisma_airs): withhold generated tool args from response-side blocks _scan_tool_calls_for_guardrail calls AIRS with is_response=False because tool_event is request-side in the AIRS schema, so AIRS returns the scanned tool arguments under prompt_masked_data. When the tool calls being scanned are the model's own output, that key holds generated content, and the _CLIENT_HIDDEN_SCAN_FIELDS default (response_masked_data, empty on this path) does not cover it. With the default mask_response_content=False the block branch then shipped the model's masked tool arguments in the 400 -- the same content channel this PR closed for response_masked_data. _build_error_detail takes an extra_hidden_fields argument so the withholding stays in one place, and the tool-call block branch passes prompt_masked_data when is_response is True. Request-side blocks are unchanged and still carry prompt_masked_data, which is what LIT-5638 asks for. Co-Authored-By: Claude Opus 5 (1M context) * style(panw_prisma_airs): apply ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Yucheng Zhu Co-authored-by: Claude Opus 5 (1M context) --- .../panw_prisma_airs/panw_prisma_airs.py | 60 +++-- .../guardrail_hooks/test_panw_prisma_airs.py | 229 ++++++++++++++++++ 2 files changed, 266 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 3765771247d..7e641814deb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -71,6 +71,14 @@ class PanwPrismaAirsHandler(CustomGuardrail): _PROVIDER_NAME = "panw_prisma_airs" + #: AIRS fields withheld from the client-visible error detail. + #: ``response_masked_data`` is the model's own generation. The block branch that builds + #: this detail is only reached when ``mask_response_content`` is False, so echoing it + #: back would hand the caller exactly the text the operator declined to deliver. + #: ``prompt_masked_data`` is deliberately NOT withheld: it is the caller's own input, + #: and it is one of the fields the ticket asks for. + _CLIENT_HIDDEN_SCAN_FIELDS: Final = frozenset({"response_masked_data"}) + def __init__( self, guardrail_name: str, @@ -632,12 +640,21 @@ class PanwPrismaAirsHandler(CustomGuardrail): choice.message.function_call.arguments = masked_text def _build_error_detail( - self, scan_result: Mapping[str, object], is_response: bool = False + self, + scan_result: Mapping[str, object], + is_response: bool = False, + also_hide: str | None = None, ) -> Mapping[str, Mapping[str, object]]: - """Build enhanced error detail with scan information.""" + """Build enhanced error detail with scan information. + + ``also_hide`` names one more scan field to withhold, for the caller that knows + its AIRS verdict carries model-generated content under a key that is normally + caller input. + """ action_type: Final = "Response" if is_response else "Prompt" code_suffix: Final = "_response_blocked" if is_response else "_blocked" - detection_key: Final = "response_detected" if is_response else "prompt_detected" + + hidden_fields: Final = self._CLIENT_HIDDEN_SCAN_FIELDS.union(() if also_hide is None else (also_hide,)) category: Final = scan_result.get("category", "unknown") default_msg: Final = f"{action_type} blocked by PANW Prisma AI Security policy (Category: {category})" @@ -653,8 +670,13 @@ class PanwPrismaAirsHandler(CustomGuardrail): }, ) - error_detail: Final[dict[str, dict[str, object]]] = { + return { "error": { + **{ + key: value + for key, value in scan_result.items() + if not key.startswith("_") and key not in hidden_fields + }, "message": error_msg, "type": "guardrail_violation", "code": f"panw_prisma_airs{code_suffix}", @@ -663,24 +685,6 @@ class PanwPrismaAirsHandler(CustomGuardrail): } } - # Add optional fields if present - optional_fields: Final = [ - "scan_id", - "report_id", - "profile_name", - "profile_id", - "tr_id", - ] - for field in optional_fields: - if scan_result.get(field): - error_detail["error"][field] = scan_result[field] - - # Add detection details - if scan_result.get(detection_key): - error_detail["error"][detection_key] = scan_result[detection_key] - - return error_detail - def _record_scan_id(self, request_data: dict[str, Any], scan_result: Mapping[str, object]) -> None: """Surface the AIRS scan id on the response, so allowed calls are auditable too.""" scan_id: Final = scan_result.get("scan_id") @@ -1481,7 +1485,17 @@ class PanwPrismaAirsHandler(CustomGuardrail): ): self._set_tool_call_arguments(tool_call, masked_text) else: - error_detail = self._build_error_detail(scan_result, is_response=is_response) + # tool_event scans are request-side in the AIRS schema, so AIRS returns + # the model's own tool arguments under prompt_masked_data. On a + # response-side block that is generated content, not caller input, and + # the class-level default only withholds response_masked_data — which is + # empty on this path. Withhold it explicitly so the 400 does not become + # the content channel this branch declined to deliver. + error_detail = self._build_error_detail( + scan_result, + is_response=is_response, + also_hide="prompt_masked_data" if is_response else None, + ) raise HTTPException(status_code=400, detail=error_detail) @staticmethod diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 17d4a3e304a..fc1465c7c14 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -5647,6 +5647,235 @@ class TestPanwAirsScanIdExposure: assert "guardrail_scan_ids" in _UNTRUSTED_METADATA_CONTROL_FIELDS assert "guardrail_scan_ids" in _UNTRUSTED_ROOT_CONTROL_FIELDS +class TestPanwAirsBlockedErrorDetailPassthrough: + """Regression tests for the full AIRS scan response on blocks. + + Before the fix, the error detail was built from a hardcoded allowlist + (scan_id, report_id, profile_name, profile_id, tr_id, prompt/response_detected), + so audit-relevant fields such as prompt_detection_details, prompt_masked_data, + source, transaction_id and session_id never reached the client. + """ + + _FULL_BLOCK_RESPONSE = { + "action": "block", + "category": "malicious", + "scan_id": "b2f0a4be-1f6f-4f9a-9f3d-4b6a9d8b1c0e", + "report_id": "R0000000000000000000", + "tr_id": "test-call-id", + "profile_id": "6f5c9f6e-2d0b-4d3f-8a1e-9b7c5d4e3f2a", + "profile_name": "test_profile", + "source": "prisma_airs", + "transaction_id": "4b8c1e2f-5a6d-4c3b-9e8f-1a2b3c4d5e6f", + "session_id": "3a2b1c0d-9e8f-4a7b-8c6d-5e4f3a2b1c0d", + "timeout": False, + "errors": [], + "prompt_detected": {"dlp": True, "injection": False, "url_cats": False}, + "prompt_detection_details": { + "dlp_report": { + "dlp_report_id": "1234567890", + "dlp_profile_name": "Sensitive Content", + "data_pattern_rule1_verdict": "MATCHED", + } + }, + "prompt_masked_data": {"data": "my ssn is XXX-XX-XXXX"}, + "response_detected": {"dlp": False, "url_cats": False}, + "response_detection_details": {}, + "response_masked_data": {}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("is_response", [False, True]) + async def test_block_returns_every_airs_field( + self, base_handler, user_api_key_dict, safe_prompt_data, is_response + ): + response = ModelResponse( + id="test_id", + choices=[ + Choices(index=0, message=Message(role="assistant", content="Test response")), + ], + model="gpt-3.5-turbo", + ) + + with patch.object( + base_handler, "_call_panw_api", return_value=copy.deepcopy(self._FULL_BLOCK_RESPONSE) + ): + with pytest.raises(HTTPException) as exc_info: + if is_response: + await base_handler.async_post_call_success_hook( + data=safe_prompt_data, + user_api_key_dict=user_api_key_dict, + response=response, + ) + else: + await base_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=None, + data=safe_prompt_data, + call_type="completion", + ) + + error = exc_info.value.detail["error"] + for field, value in self._FULL_BLOCK_RESPONSE.items(): + if field == "category": + continue + if field in PanwPrismaAirsHandler._CLIENT_HIDDEN_SCAN_FIELDS: + # Withheld on purpose, covered by TestPanwAirsErrorDetailWithheldFields + continue + assert error[field] == value, f"{field} missing or altered in blocked-request error" + + assert error["category"] == "malicious" + assert error["type"] == "guardrail_violation" + assert error["guardrail"] == "test_panw_airs" + assert error["code"] == ("panw_prisma_airs_response_blocked" if is_response else "panw_prisma_airs_blocked") + assert "PANW Prisma AI Security policy" in error["message"] + + def test_internal_control_flags_are_not_leaked(self, base_handler): + detail = base_handler._build_error_detail( + { + "action": "block", + "category": "malicious", + "scan_id": "scan-1", + "_always_block": True, + "_is_transient": True, + } + ) + + assert "_always_block" not in detail["error"] + assert "_is_transient" not in detail["error"] + assert detail["error"]["scan_id"] == "scan-1" + + +class TestPanwAirsErrorDetailWithheldFields: + """The blocked-request passthrough must not become a content channel. + + ``response_masked_data`` is the model's own generation. The block branch is only + reached when ``mask_response_content`` is False, so echoing it back would hand the + caller exactly the text the operator declined to deliver. ``error`` is AIRS's own + message about the operator's Strata Cloud Manager profile configuration. + + ``prompt_masked_data`` is deliberately NOT withheld by default: it is the caller's + own input, and it is one of the fields LIT-5638 asks for. The one exception is the + response-side tool-call path, covered by + ``TestPanwAirsToolCallBlockWithholdsGeneratedArgs`` below — tool_event scans are + request-side in the AIRS schema, so there the key holds model output instead. + """ + + @pytest.mark.parametrize("is_response", [False, True]) + def test_response_masked_data_never_reaches_client(self, base_handler, is_response): + detail = base_handler._build_error_detail( + { + "action": "block", + "category": "sensitive_data", + "scan_id": "scan-1", + "response_detected": {"dlp": True}, + "response_masked_data": {"data": "routing number XXXXXXXXXX"}, + "prompt_masked_data": {"data": "my ssn is XXX-XX-XXXX"}, + "prompt_detection_details": {"dlp_report": {"dlp_report_id": "1"}}, + }, + is_response=is_response, + ) + error = detail["error"] + + assert "response_masked_data" not in error + assert "routing number" not in str(error) + + # The audit fields LIT-5638 asks for still come through untouched. + assert error["scan_id"] == "scan-1" + assert error["response_detected"] == {"dlp": True} + assert error["prompt_masked_data"] == {"data": "my ssn is XXX-XX-XXXX"} + assert error["prompt_detection_details"] == {"dlp_report": {"dlp_report_id": "1"}} + + def test_upstream_airs_error_field_still_passes_through(self, base_handler): + """A 2xx AIRS body can carry its own ``error`` (see _call_panw_api's + profile-misconfiguration branch, which only logs and then blocks). It is + diagnostic rather than content, so it stays in the passthrough.""" + detail = base_handler._build_error_detail( + { + "action": "block", + "category": "malicious", + "scan_id": "scan-2", + "error": "profile not found", + } + ) + + assert detail["error"]["error"] == "profile not found" + assert detail["error"]["scan_id"] == "scan-2" + + +class TestPanwAirsToolCallBlockWithholdsGeneratedArgs: + """A response-side tool-call block must not ship the model's tool arguments. + + ``_scan_tool_calls_for_guardrail`` calls AIRS with ``is_response=False`` because + tool_event is request-side in the AIRS schema, so AIRS returns the scanned tool + arguments under ``prompt_masked_data``. When the tool calls being scanned are the + model's own output, that key holds generated content, and the class-level + ``_CLIENT_HIDDEN_SCAN_FIELDS`` default (``response_masked_data``, empty on this + path) does not cover it. + """ + + MASKED_ARGS = '{"to_account": "XXXXXXXXXX", "amount": 5000}' + + SCAN_RESULT = { + "action": "block", + "category": "sensitive_data", + "scan_id": "scan-tool-1", + "prompt_detected": {"dlp": True}, + "prompt_masked_data": {"data": MASKED_ARGS}, + "response_masked_data": {}, + } + + @staticmethod + def _tool_call(): + return ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function( + name="transfer_funds", + arguments='{"to_account": "ACME-VENDOR-001", "amount": 5000}', + ), + ) + + async def _block(self, handler, is_response): + with patch.object(handler, "_call_panw_api", new_callable=AsyncMock) as mock_api: + mock_api.return_value = dict(self.SCAN_RESULT) + with pytest.raises(HTTPException) as exc_info: + await handler._scan_tool_calls_for_guardrail( + tool_calls=[self._tool_call()], + is_response=is_response, + metadata={}, + call_id="test-call-id", + request_data={"metadata": {}}, + start_time=datetime.now(), + ) + return exc_info.value + + @pytest.mark.asyncio + async def test_response_side_block_withholds_generated_tool_args(self): + handler = make_handler(mask_response_content=False) + # The block branch is only reached with masking off; guard the premise. + assert handler.mask_response_content is False + + exc = await self._block(handler, is_response=True) + error = exc.detail["error"] + + assert exc.status_code == 400 + assert "prompt_masked_data" not in error + assert self.MASKED_ARGS not in str(error) + + # The audit fields LIT-5638 asks for are unaffected. + assert error["scan_id"] == "scan-tool-1" + assert error["prompt_detected"] == {"dlp": True} + + @pytest.mark.asyncio + async def test_request_side_block_still_returns_masked_tool_args(self): + """Caller-supplied tool arguments stay in the verdict — that is the ticket's ask.""" + handler = make_handler(mask_request_content=False) + + exc = await self._block(handler, is_response=False) + error = exc.detail["error"] + + assert error["prompt_masked_data"] == {"data": self.MASKED_ARGS} + assert error["scan_id"] == "scan-tool-1" if __name__ == "__main__":