From 361170f9d4ec67d92eec187710f93ec81de9f729 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 23:19:04 +0000 Subject: [PATCH 01/61] feat(organization): expose PATCH /v2/organization/{organization_id} in the OpenAPI spec Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/organization_endpoints.py | 1 - .../test_organization_endpoints.py | 14 ++++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 5e38a016099..35a1380a619 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -764,7 +764,6 @@ async def handle_update_object_permission( tags=["organization management"], dependencies=[Depends(user_api_key_auth)], response_model=LiteLLM_OrganizationTableWithMembers, - include_in_schema=False, ) async def update_organization_v2( organization_id: str, diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index e2d89a660c2..5289f4f2d8f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -1063,3 +1063,17 @@ async def test_find_member_if_email_missing_row_raises_documented_400(): "non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead." ) } + + +def test_v2_update_organization_is_in_openapi_schema(): + """PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec.""" + from fastapi import FastAPI + + from litellm.proxy.management_endpoints.organization_endpoints import router + + app = FastAPI() + app.include_router(router) + + v2_path = app.openapi()["paths"]["/v2/organization/{organization_id}"] + assert v2_path["patch"]["tags"] == ["organization management"] + assert "OrganizationUpdateRequestV2" in json.dumps(v2_path["patch"]["requestBody"]) From 9ba6cab889c01edd360cac5a4b38e6a942bcafbf Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 23:59:58 +0000 Subject: [PATCH 02/61] fix(ui): make Admin UI table pagination honor the selected page size All Models now pushes the model group, access group and wildcard filters into /v2/model/info (new optional access_group and wildcard_only params) so the server total_count matches the rendered rows. Request Logs defaults to 25, uses the shared page size options and counts rendered rows in the footer. Deleted Teams gets the shared DataTable server pagination footer instead of a hard-coded page size of 100. Per-user usage and the remaining unbounded list tables get paginationMode so the size selector renders. Resolves LIT-4738 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 36 +++++++- .../proxy_server/test_routes_model_info.py | 89 +++++++++++++++++++ .../agents/_components/AgentsTable.tsx | 1 + .../_components/guardrail_table.tsx | 1 + .../app/(dashboard)/hooks/models/useModels.ts | 6 ++ .../(dashboard)/hooks/teams/useTeams.test.ts | 24 ++++- .../app/(dashboard)/hooks/teams/useTeams.ts | 21 +++-- .../_components/MCPToolsetsTab.tsx | 1 + .../components/AllModelsTab.test.tsx | 49 ++++++++-- .../components/AllModelsTab.tsx | 32 ++----- .../panels/AccessGroupBudgetsPanel.tsx | 1 + .../_components/OrganizationsTable.test.tsx | 17 ++++ .../_components/OrganizationsTable.tsx | 1 + .../policies/_components/AttachmentTable.tsx | 1 + .../policies/_components/PolicyTable.tsx | 1 + .../prompts/_components/PromptTable.tsx | 1 + .../_components/SearchToolTable.tsx | 1 + .../skills/_components/PluginTable.tsx | 1 + .../tag-management/_components/TagTable.tsx | 1 + .../_components/IndexesTable.tsx | 1 + .../_components/VectorStoreTable.tsx | 1 + .../src/components/AIHub/ModelHubTable.tsx | 3 + .../components/AIHub/SkillHubDashboard.tsx | 1 + .../DeletedTeamsPage.test.tsx | 48 +++++++++- .../DeletedTeamsPage/DeletedTeamsPage.tsx | 17 +++- .../DeletedTeamsTable.test.tsx | 32 ++++++- .../DeletedTeamsTable/DeletedTeamsTable.tsx | 17 +++- .../PassThroughEndpointsTable.tsx | 1 + .../components/model_add/CredentialsTable.tsx | 1 + .../src/components/networking.tsx | 8 ++ .../src/components/per_user_usage.test.tsx | 19 ++++ .../src/components/per_user_usage.tsx | 53 +++-------- .../src/components/public_model_hub.tsx | 3 + .../routing_groups/RoutingGroupsTable.tsx | 1 + .../components/team/AvailableTeamsTable.tsx | 1 + .../view_logs/RequestLogsPanel.test.tsx | 51 ++++++++++- .../components/view_logs/RequestLogsPanel.tsx | 10 ++- .../components/view_logs/RequestLogsTable.tsx | 2 - .../src/components/view_logs/constants.ts | 3 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 ++ 40 files changed, 462 insertions(+), 102 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c43cc510990..1d3455615fd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13474,6 +13474,27 @@ def _is_auto_router_model(model: Mapping[str, object]) -> bool: return isinstance(litellm_model, str) and litellm_model.startswith("auto_router/") +def _model_in_access_group(model: Mapping[str, object], access_group: str) -> bool: + model_info: Final = model.get("model_info") + if not isinstance(model_info, Mapping): + return False + access_groups: Final = model_info.get("access_groups") + return isinstance(access_groups, (list, tuple)) and access_group in access_groups + + +def _matches_model_info_filters( + model: Mapping[str, object], + exclude_auto_routers: bool | None, + access_group: str | None, + wildcard_only: bool | None, +) -> bool: + if exclude_auto_routers is True and _is_auto_router_model(model): + return False + if isinstance(access_group, str) and not _model_in_access_group(model, access_group): + return False + return wildcard_only is not True or "*" in str(model.get("model_name") or "") + + def _paginate_models_response( all_models: list[dict[str, Any]], page: int, @@ -13784,6 +13805,14 @@ async def model_info_v2( "existing callers are unaffected" ), ), + access_group: str | None = fastapi.Query( + None, + description="Only return deployments whose `model_info.access_groups` contains this access group", + ), + wildcard_only: bool | None = fastapi.Query( + False, + description="Only return wildcard deployments, i.e. those whose `model_name` contains `*`", + ), ): """ Paginated model metadata for proxy deployments (pricing, provider, team access). @@ -13801,6 +13830,8 @@ async def model_info_v2( modelId: Return a single deployment by LiteLLM model id. teamId: Filter to models with direct access or team membership for this team id. sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. + access_group: Only return deployments in this model access group. + wildcard_only: Only return deployments whose `model_name` contains `*`. Example request: ``` @@ -13954,8 +13985,9 @@ async def model_info_v2( # `is True` because direct-call tests bypass FastAPI, so the Query default arrives as a # truthy sentinel object rather than False. - if exclude_auto_routers is True: - all_models = [m for m in all_models if not _is_auto_router_model(m)] + all_models = [ + m for m in all_models if _matches_model_info_filters(m, exclude_auto_routers, access_group, wildcard_only) + ] # Update total count to include agents search_total_count = len(all_models) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index cb38e7edbe2..4c141bcf698 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -452,3 +452,92 @@ async def test_model_info_v2_query_sentinel_does_not_filter(monkeypatch, mixed_a ) assert "tri-tier-router" in [m["model_name"] for m in resp["data"]] + + +# --------------------------------------------------------------------------- +# GET /v2/model/info?access_group / ?wildcard_only +# --------------------------------------------------------------------------- + + +@pytest.fixture +def access_group_router(monkeypatch): + """Router with one sales-team deployment, one wildcard sales-team deployment and one ungrouped one.""" + model_list = [ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "sales-1", "db_model": False, "access_groups": ["sales-team"]}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "sales-wildcard", "db_model": False, "access_groups": ["sales-team", "eng"]}, + }, + { + "model_name": "claude-opus", + "litellm_params": {"model": "anthropic/claude-opus-4-6"}, + "model_info": {"id": "plain-1", "db_model": False}, + }, + ] + from unittest.mock import AsyncMock + + router = MagicMock() + router.model_list = model_list + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", model_list) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server.proxy_config, "get_config", AsyncMock(return_value={})) + monkeypatch.setattr( + proxy_server, + "_apply_search_filter_to_models", + AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))), + ) + monkeypatch.setattr(proxy_server, "_enrich_model_info_with_litellm_data", lambda model, **kw: model) + + import litellm.proxy.agent_endpoints.model_list_helpers as mlh + + monkeypatch.setattr(mlh, "append_agents_to_model_info", AsyncMock(side_effect=lambda models, **kw: models)) + yield router + + +def test_v2_model_info_without_new_filters_returns_everything(client, auth_as, access_group_router): + with auth_as(): + response = client.get("/v2/model/info") + payload = response.json() + assert payload["total_count"] == 3 + assert len(payload["data"]) == 3 + + +def test_v2_model_info_access_group_filters_rows_and_total(client, auth_as, access_group_router): + """The table pages off total_count, so the filter must shrink the total, not only the page.""" + with auth_as(): + response = client.get("/v2/model/info", params={"access_group": "sales-team"}) + payload = response.json() + assert _model_names(payload) == ["gpt-4o-mini", "openai/*"] + assert payload["total_count"] == 2 + + +def test_v2_model_info_unknown_access_group_is_empty(client, auth_as, access_group_router): + with auth_as(): + response = client.get("/v2/model/info", params={"access_group": "nobody"}) + payload = response.json() + assert payload["data"] == [] + assert payload["total_count"] == 0 + + +def test_v2_model_info_wildcard_only_filters_rows_and_total(client, auth_as, access_group_router): + with auth_as(): + response = client.get("/v2/model/info", params={"wildcard_only": "true"}) + payload = response.json() + assert _model_names(payload) == ["openai/*"] + assert payload["total_count"] == 1 + + +def test_v2_model_info_access_group_paginates_over_the_filtered_set(client, auth_as, access_group_router): + with auth_as(): + response = client.get("/v2/model/info", params={"access_group": "sales-team", "page": 2, "size": 1}) + payload = response.json() + assert _model_names(payload) == ["openai/*"] + assert payload["total_count"] == 2 + assert payload["total_pages"] == 2 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx index aceb07e2e9a..d737a9250eb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -72,6 +72,7 @@ const AgentsTable: React.FC = ({ return ( agent.agent_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx index e6a14b2b2f4..bbab01e346d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx @@ -46,6 +46,7 @@ const GuardrailTable: React.FC = ({ return ( guardrail.guardrail_id || String(index)} sortingMode="client" 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 a9f7c54698a..b3a783a71dc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -39,6 +39,8 @@ export const useModelsInfo = ( sortOrder?: string, excludeAutoRouters: boolean = false, modelName?: string, + accessGroup?: string, + wildcardOnly: boolean = false, ) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ @@ -57,6 +59,8 @@ export const useModelsInfo = ( // Part of the key: callers that exclude auto-routers must not share a cache entry // with callers that keep them. ...(excludeAutoRouters && { excludeAutoRouters: "true" }), + ...(accessGroup && { accessGroup }), + ...(wildcardOnly && { wildcardOnly: "true" }), }, }), queryFn: async () => @@ -73,6 +77,8 @@ export const useModelsInfo = ( sortOrder, excludeAutoRouters, modelName, + accessGroup, + wildcardOnly, ), enabled: Boolean(accessToken && userId && userRole), }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts index fa3f15124cf..eccd8a80748 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -671,7 +671,7 @@ describe("useDeletedTeams", () => { it("should return deleted teams data when query is successful", async () => { (global.fetch as any).mockResolvedValue({ ok: true, - json: async () => ({ teams: mockDeletedTeams }), + json: async () => ({ teams: mockDeletedTeams, total: 2, page: 1, page_size: 10, total_pages: 1 }), }); const { result } = renderHook(() => useDeletedTeams(1, 10, {}), { wrapper }); @@ -684,10 +684,26 @@ describe("useDeletedTeams", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual(mockDeletedTeams); + expect(result.current.data).toEqual({ teams: mockDeletedTeams, total: 2 }); expect(result.current.error).toBeNull(); }); + it("should keep the server total so the table can paginate beyond the current page", async () => { + (global.fetch as any).mockResolvedValue({ + ok: true, + json: async () => ({ teams: mockDeletedTeams, total: 137, page: 1, page_size: 2, total_pages: 69 }), + }); + + const { result } = renderHook(() => useDeletedTeams(1, 2, {}), { wrapper }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data?.total).toBe(137); + expect((global.fetch as any).mock.calls[0][0]).toContain("page_size=2"); + }); + it("should handle error when API call fails", async () => { (global.fetch as any).mockResolvedValue({ ok: false, @@ -744,7 +760,7 @@ describe("useDeletedTeams", () => { rerender({ page: 2 }); - expect(result.current.data).toEqual(mockDeletedTeams); + expect(result.current.data?.teams).toEqual(mockDeletedTeams); }); it("should pass options to API call", async () => { @@ -785,7 +801,7 @@ describe("useDeletedTeams", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual(mockDeletedTeams); + expect(result.current.data).toEqual({ teams: mockDeletedTeams, total: 2 }); expect(result.current.error).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index e209a1d7273..14e95bcd543 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -20,6 +20,11 @@ export interface DeletedTeam extends Team { deleted_by: string; } +export interface DeletedTeamsResponse { + teams: DeletedTeam[]; + total: number; +} + export interface TeamListCallOptions { organizationID?: string | null; teamID?: string | null; @@ -209,7 +214,7 @@ const deletedTeamListCall = async ( page: number, pageSize: number, options: TeamListCallOptions = {}, -) => { +): Promise => { /** * Get deleted teams from proxy */ @@ -251,14 +256,12 @@ const deletedTeamListCall = async ( throw new Error(errorMessage); } - const data = await response.json(); + const data: DeletedTeam[] | (Partial & { teams: DeletedTeam[] }) = await response.json(); - // Extract teams array from response if it's wrapped in a response object - // Otherwise return the data directly if it's already an array - if (data && typeof data === "object" && "teams" in data) { - return data.teams as DeletedTeam[]; + if (Array.isArray(data)) { + return { teams: data, total: data.length }; } - return data as DeletedTeam[]; + return { teams: data.teams, total: data.total ?? data.teams.length }; } catch (error) { console.error("Failed to list deleted teams:", error); throw error; @@ -270,10 +273,10 @@ export const useDeletedTeams = ( page: number, pageSize: number, options: TeamListCallOptions = {}, -): UseQueryResult => { +): UseQueryResult => { const { accessToken } = useAuthorized(); - return useQuery({ + return useQuery({ queryKey: deletedTeamKeys.list({ page, limit: pageSize, ...options }), queryFn: async () => await deletedTeamListCall(accessToken!, page, pageSize, options), enabled: Boolean(accessToken), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx index c637655d665..60a1da40d08 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx @@ -451,6 +451,7 @@ export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) { toolset.toolset_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 65faa85e29e..7e47be3f5d1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -34,6 +34,8 @@ interface ModelsInfoArgs { sortBy?: string; sortOrder?: string; modelName?: string; + accessGroup?: string; + wildcardOnly?: boolean; } const modelsInfoCalls: ModelsInfoArgs[] = []; @@ -50,12 +52,24 @@ type UseModelsInfoArgs = [ sortOrder?: string, excludeAutoRouters?: boolean, modelName?: string, + accessGroup?: string, + wildcardOnly?: boolean, ]; vi.mock("../../hooks/models/useModels", () => ({ useModelsInfo: (...args: UseModelsInfoArgs) => { - const [page, size, search, , teamId, sortBy, sortOrder, , modelName] = args; - const call: ModelsInfoArgs = { page, size, search, teamId, sortBy, sortOrder, modelName }; + const [page, size, search, , teamId, sortBy, sortOrder, , modelName, accessGroup, wildcardOnly] = args; + const call: ModelsInfoArgs = { + page, + size, + search, + teamId, + sortBy, + sortOrder, + modelName, + accessGroup, + wildcardOnly, + }; modelsInfoCalls.push(call); return { ...modelsInfoResult, refetch: mockRefetch }; }, @@ -254,13 +268,38 @@ describe("AllModelsTab", () => { }); }); - it("filters the fetched page down to the selected model group", () => { - setModelsInfo([makeRow(), { ...makeRow(), model_name: "claude-opus" }], 2); + it("renders every row the server returned for the selected model group so rows match the footer total", () => { + setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "claude-opus" }], 2); render(); const table = screen.getByRole("table"); expect(within(table).getByText("claude-opus")).toBeInTheDocument(); - expect(within(table).queryByText("gpt-4")).not.toBeInTheDocument(); + expect(within(table).getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-2 of 2"); + }); + + it("asks the server for wildcard deployments instead of hiding rows client-side", () => { + setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "openai/*" }], 2); + render(); + + expect(lastModelsInfoCall().wildcardOnly).toBe(true); + expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-2 of 2"); + }); + + it("asks the server for the selected access group instead of hiding rows client-side", async () => { + const user = userEvent.setup(); + render(); + expect(lastModelsInfoCall().wildcardOnly).toBe(false); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(await screen.findByPlaceholderText("Filter by Model Access Group")); + await user.click(await screen.findByRole("option", { name: "sales-team" })); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastModelsInfoCall().accessGroup).toBe("sales-team")); + expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-1 of 1"); }); it("asks the server for the exact selected model group so deployments beyond the first page are found", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index be2cf22d71a..3b4058a28fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -86,6 +86,11 @@ const AllModelsTab = ({ selectedModelGroup !== ALL_MODEL_GROUPS_VALUE && selectedModelGroup !== WILDCARD_MODEL_GROUP_VALUE; const modelNameForQuery = isConcreteModelGroup ? selectedModelGroup ?? undefined : undefined; + const accessGroupForQuery = + selectedModelAccessGroupFilter && selectedModelAccessGroupFilter !== ALL_MODEL_GROUPS_VALUE + ? selectedModelAccessGroupFilter + : undefined; + const wildcardOnlyForQuery = selectedModelGroup === WILDCARD_MODEL_GROUP_VALUE; const sortBy = useMemo(() => { if (sorting.length === 0) return undefined; @@ -114,6 +119,8 @@ const AllModelsTab = ({ // lists and manages them. Excluded server-side so total_count stays honest. true, modelNameForQuery, + accessGroupForQuery, + wildcardOnlyForQuery, ); const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; @@ -129,32 +136,11 @@ const AllModelsTab = ({ [modelCostMapData], ); - const modelData = useMemo(() => { + const modelData = useMemo<{ data: ModelData[] }>(() => { if (!rawModelData) return { data: [] }; return transformModelData(rawModelData, getProviderFromModel); }, [rawModelData, getProviderFromModel]); - const filteredData = useMemo(() => { - if (!modelData || !modelData.data || modelData.data.length === 0) { - return []; - } - - return modelData.data.filter((model: ModelData) => { - const modelNameMatch = - selectedModelGroup === ALL_MODEL_GROUPS_VALUE || - model.model_name === selectedModelGroup || - !selectedModelGroup || - (selectedModelGroup === WILDCARD_MODEL_GROUP_VALUE && model.model_name?.includes("*")); - - const accessGroupMatch = - selectedModelAccessGroupFilter === ALL_MODEL_GROUPS_VALUE || - model.model_info["access_groups"]?.includes(selectedModelAccessGroupFilter ?? "") || - !selectedModelAccessGroupFilter; - - return modelNameMatch && accessGroupMatch; - }); - }, [modelData, selectedModelGroup, selectedModelAccessGroupFilter]); - const columnFilters = useMemo( () => [ @@ -270,7 +256,7 @@ const AllModelsTab = ({
group.access_group} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx index 1ac33a27186..4bf465b847b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -192,6 +192,23 @@ describe("OrganizationsTable", () => { expect(screen.queryByText("ShouldNotShow")).not.toBeInTheDocument(); }); + it("pages long lists client-side with the shared size selector and footer", async () => { + const user = userEvent.setup(); + const organizations = Array.from({ length: 30 }, (_, index) => + makeOrganization({ organization_id: `org-${index}`, organization_alias: `Org ${index}` }), + ); + render(); + + expect(screen.getAllByRole("row")).toHaveLength(26); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30"); + + await user.click(screen.getByTestId("pagination-page-size")); + await user.click(await screen.findByRole("option", { name: "50" })); + + expect(screen.getAllByRole("row")).toHaveLength(31); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-30 of 30"); + }); + it("uses a search-aware empty state", () => { const { rerender } = render(); expect(screen.getByText("No organizations yet")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx index 8e68a57d2f7..dbf516d75ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx @@ -59,6 +59,7 @@ const OrganizationsTable: React.FC = ({ return ( organization.organization_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.tsx index bd8458e6f96..a432a53bca4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.tsx @@ -50,6 +50,7 @@ const AttachmentTable: React.FC = ({ return ( row.attachment_id} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx index 3405ac6b6bb..d78ec28c486 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx @@ -71,6 +71,7 @@ const PolicyTable: React.FC = ({ return ( `${row.primaryPolicy.definition_location ?? "db"}:${row.policy_name}`} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx index c766042ac44..e810c3622d1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx @@ -73,6 +73,7 @@ const PromptTable: React.FC = ({ return ( prompt.prompt_id ? `${prompt.prompt_id}::${prompt.environment || "development"}` : String(index) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx index 70fc6a376df..f60f4f3d1da 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx @@ -50,6 +50,7 @@ const SearchToolTable: React.FC = ({ return ( searchToolKey(tool) || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx index c581b0dfdeb..1b1ccb0932a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx @@ -42,6 +42,7 @@ const PluginTable: React.FC = ({ pluginsList, isLoading, onDel return ( plugin.id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx index 488190a0fdf..076166ac827 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx @@ -39,6 +39,7 @@ const TagTable: React.FC = ({ data, onEdit, onDelete, onSelectTag return ( tag.name || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx index 927fd48acb6..a0b0c02f99d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx @@ -46,6 +46,7 @@ const IndexesTable: React.FC = ({ return ( row.id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx index 2f8508dc7c6..32e7bc2324d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx @@ -41,6 +41,7 @@ const VectorStoreTable: React.FC = ({ data, onView, onEdi return ( vectorStore.vector_store_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 475e3dcd70b..5f6f26bd16c 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -474,6 +474,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Model Table */} model.model_group || String(index)} sortingMode="client" @@ -540,6 +541,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Agent Table */} agent.agent_id || agent.name || String(index)} sortingMode="client" @@ -581,6 +583,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* MCP Server Table */} server.server_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx b/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx index 992ef49742d..9cede3b4497 100644 --- a/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx @@ -162,6 +162,7 @@ const SkillHubDashboard: React.FC = ({
skill.id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx index 6bf5d1caf61..952d8764463 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx @@ -1,4 +1,5 @@ -import { screen } from "@testing-library/react"; +import { fireEvent, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { vi, it, expect, beforeEach, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import DeletedTeamsPage from "./DeletedTeamsPage"; @@ -31,7 +32,7 @@ beforeEach(() => { vi.clearAllMocks(); mockUseDeletedTeams.mockReturnValue({ - data: [mockDeletedTeam], + data: { teams: [mockDeletedTeam], total: 1 }, isLoading: false, } as unknown as ReturnType); }); @@ -42,6 +43,49 @@ it("should render DeletedTeamsPage component", () => { expect(screen.getByText("Test Team")).toBeInTheDocument(); }); +it("requests the first page of 25 deleted teams and shows the server total in the footer", () => { + mockUseDeletedTeams.mockReturnValue({ + data: { teams: [mockDeletedTeam], total: 137 }, + isLoading: false, + } as unknown as ReturnType); + + renderWithProviders(); + + expect(mockUseDeletedTeams).toHaveBeenLastCalledWith(1, 25); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 137"); + expect(screen.getByTestId("pagination-next")).toBeEnabled(); +}); + +it("requests the next page from the server when Next is clicked", () => { + mockUseDeletedTeams.mockReturnValue({ + data: { teams: [mockDeletedTeam], total: 137 }, + isLoading: false, + } as unknown as ReturnType); + + renderWithProviders(); + fireEvent.click(screen.getByTestId("pagination-next")); + + expect(mockUseDeletedTeams).toHaveBeenLastCalledWith(2, 25); +}); + +it("offers the shared page sizes and refetches with the selected one", async () => { + const user = userEvent.setup(); + mockUseDeletedTeams.mockReturnValue({ + data: { teams: [mockDeletedTeam], total: 137 }, + isLoading: false, + } as unknown as ReturnType); + + renderWithProviders(); + await user.click(screen.getByTestId("pagination-page-size")); + + const options = await screen.findAllByRole("option"); + expect(options.map((option) => option.textContent)).toEqual(["25", "50", "100"]); + + await user.click(screen.getByRole("option", { name: "100" })); + + expect(mockUseDeletedTeams).toHaveBeenLastCalledWith(1, 100); +}); + it("should show the enterprise notice for a non-premium user", () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx index eab150d6ab5..8c3aac2cac7 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx @@ -1,13 +1,20 @@ "use client"; +import { PaginationState } from "@tanstack/react-table"; import { Info } from "lucide-react"; +import { useState } from "react"; import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { DEFAULT_PAGE_SIZE_OPTIONS } from "@/components/shared/DataTable"; import { useDeletedTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { DeletedTeamsTable } from "./DeletedTeamsTable/DeletedTeamsTable"; export default function DeletedTeamsPage() { const { premiumUser } = useAuthorized(); - const { data: teamsData, isLoading } = useDeletedTeams(1, 100); + const [pagination, setPagination] = useState({ + pageIndex: 0, + pageSize: DEFAULT_PAGE_SIZE_OPTIONS[0], + }); + const { data: teamsData, isLoading } = useDeletedTeams(pagination.pageIndex + 1, pagination.pageSize); return (
@@ -20,7 +27,13 @@ export default function DeletedTeamsPage() { )} - +
); } diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx index c0cc5a342a8..e166f6b0d1b 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx @@ -22,12 +22,19 @@ const makeDeletedTeam = (overrides: Partial = {}): DeletedTeam => ( ...overrides, }); +const paginationProps = { + pagination: { pageIndex: 0, pageSize: 25 }, + onPaginationChange: vi.fn(), +}; + beforeEach(() => { vi.clearAllMocks(); }); it("should display team information", () => { - renderWithProviders(); + renderWithProviders( + , + ); expect(screen.getByText("Test Team")).toBeInTheDocument(); expect(screen.getByText("team-1")).toBeInTheDocument(); @@ -39,7 +46,7 @@ it("should sort teams by deleted_at descending by default", () => { makeDeletedTeam({ team_id: "team-old", team_alias: "older-team", deleted_at: "2024-01-01T10:00:00Z" }), makeDeletedTeam({ team_id: "team-new", team_alias: "newer-team", deleted_at: "2024-06-01T10:00:00Z" }), ]; - renderWithProviders(); + renderWithProviders(); const rows = screen.getAllByRole("row").slice(1); expect(within(rows[0]).getByText("newer-team")).toBeInTheDocument(); @@ -47,13 +54,30 @@ it("should sort teams by deleted_at descending by default", () => { }); it("should show skeleton rows when loading", () => { - renderWithProviders(); + renderWithProviders(); expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0); }); it("should show the empty state when there are no deleted teams", () => { - renderWithProviders(); + renderWithProviders(); expect(screen.getByText("No deleted teams found")).toBeInTheDocument(); }); + +it("renders the shared pagination footer with the server row count", () => { + renderWithProviders( + , + ); + + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 101-137 of 137"); + expect(screen.getByTestId("pagination-page-size")).toHaveTextContent("50"); + expect(screen.getByTestId("pagination-prev")).toBeEnabled(); + expect(screen.getByTestId("pagination-next")).toBeDisabled(); +}); diff --git a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx index 9578a52453f..c7e759754b8 100644 --- a/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx @@ -1,6 +1,6 @@ "use client"; -import { SortingState } from "@tanstack/react-table"; +import { OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import { Inbox } from "lucide-react"; import { useMemo, useState } from "react"; @@ -12,6 +12,9 @@ import { getDeletedTeamsTableColumns } from "./DeletedTeamsTableColumns"; interface DeletedTeamsTableProps { teams: DeletedTeam[]; isLoading: boolean; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + rowCount: number; } const DEFAULT_SORTING: SortingState = [{ id: "deleted_at", desc: true }]; @@ -28,7 +31,13 @@ function EmptyState() { ); } -export function DeletedTeamsTable({ teams, isLoading }: DeletedTeamsTableProps) { +export function DeletedTeamsTable({ + teams, + isLoading, + pagination, + onPaginationChange, + rowCount, +}: DeletedTeamsTableProps) { const [sorting, setSorting] = useState(DEFAULT_SORTING); const columns = useMemo(() => getDeletedTeamsTableColumns(), []); @@ -41,6 +50,10 @@ export function DeletedTeamsTable({ teams, isLoading }: DeletedTeamsTableProps) sortingMode="client" sorting={sorting} onSortingChange={setSorting} + paginationMode="server" + pagination={pagination} + onPaginationChange={onPaginationChange} + rowCount={rowCount} isLoading={isLoading} loadingMessage="Loading deleted teams…" noDataMessage={} diff --git a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.tsx b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.tsx index 754e7ff68dd..35f0bd4bc62 100644 --- a/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.tsx +++ b/ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.tsx @@ -41,6 +41,7 @@ export function PassThroughEndpointsTable({ return ( endpoint.id || endpoint.path || String(index)} isLoading={isLoading} diff --git a/ui/litellm-dashboard/src/components/model_add/CredentialsTable.tsx b/ui/litellm-dashboard/src/components/model_add/CredentialsTable.tsx index 33d63e87a5b..835cd57ae92 100644 --- a/ui/litellm-dashboard/src/components/model_add/CredentialsTable.tsx +++ b/ui/litellm-dashboard/src/components/model_add/CredentialsTable.tsx @@ -48,6 +48,7 @@ const CredentialsTable: React.FC = ({ return ( credential.credential_name || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 1384679a88a..8787b8111c6 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1692,6 +1692,8 @@ export const modelInfoCall = async ( sortOrder?: string, excludeAutoRouters?: boolean, modelName?: string, + accessGroup?: string, + wildcardOnly?: boolean, ) => { /** * Get all models on proxy @@ -1723,6 +1725,12 @@ export const modelInfoCall = async ( if (excludeAutoRouters) { params.append("exclude_auto_routers", "true"); } + if (accessGroup && accessGroup.trim()) { + params.append("access_group", accessGroup.trim()); + } + if (wildcardOnly) { + params.append("wildcard_only", "true"); + } if (params.toString()) { url += `?${params.toString()}`; } diff --git a/ui/litellm-dashboard/src/components/per_user_usage.test.tsx b/ui/litellm-dashboard/src/components/per_user_usage.test.tsx index 9cd199d786c..5cd0591bb15 100644 --- a/ui/litellm-dashboard/src/components/per_user_usage.test.tsx +++ b/ui/litellm-dashboard/src/components/per_user_usage.test.tsx @@ -78,6 +78,25 @@ describe("PerUserUsage", () => { }); }); + it("shows every fetched row with a footer that matches the server total and page size", async () => { + const results = Array.from({ length: 25 }, (_, index) => userRow(`user-${index}`, "curl/8.0", index)); + mockPerUserAnalyticsCall.mockResolvedValue({ ...mockResponse, results, total_count: 60, total_pages: 3 }); + render(); + + await waitFor(() => { + expect(screen.getByText("user-24")).toBeInTheDocument(); + }); + + expect(mockPerUserAnalyticsCall).toHaveBeenLastCalledWith("test-token", 1, 25, undefined); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 60"); + + fireEvent.click(screen.getByTestId("pagination-next")); + + await waitFor(() => { + expect(mockPerUserAnalyticsCall).toHaveBeenLastCalledWith("test-token", 2, 25, undefined); + }); + }); + it("keeps both tab panels mounted so switching tabs does not reset their state", async () => { render(); diff --git a/ui/litellm-dashboard/src/components/per_user_usage.tsx b/ui/litellm-dashboard/src/components/per_user_usage.tsx index 6f29077de84..5bdf02e61ca 100644 --- a/ui/litellm-dashboard/src/components/per_user_usage.tsx +++ b/ui/litellm-dashboard/src/components/per_user_usage.tsx @@ -1,8 +1,7 @@ import React, { useState, useEffect } from "react"; -import type { ColumnDef } from "@tanstack/react-table"; +import type { ColumnDef, PaginationState } from "@tanstack/react-table"; import { BarChart } from "@/components/shared/charts"; -import { DataTable } from "@/components/shared/DataTable"; -import { Button } from "@/components/ui/button"; +import { DataTable, DEFAULT_PAGE_SIZE_OPTIONS } from "@/components/shared/DataTable"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { perUserAnalyticsCall } from "./networking"; @@ -42,7 +41,10 @@ const PerUserUsage: React.FC = ({ accessToken, selectedTags, total_pages: 0, }); - const [currentPage, setCurrentPage] = useState(1); + const [pagination, setPagination] = useState({ + pageIndex: 0, + pageSize: DEFAULT_PAGE_SIZE_OPTIONS[0], + }); const fetchPerUserData = async () => { if (!accessToken) return; @@ -50,8 +52,8 @@ const PerUserUsage: React.FC = ({ accessToken, selectedTags, try { const response = await perUserAnalyticsCall( accessToken, - currentPage, - 50, + pagination.pageIndex + 1, + pagination.pageSize, selectedTags.length > 0 ? selectedTags : undefined, ); setPerUserData(response); @@ -62,19 +64,7 @@ const PerUserUsage: React.FC = ({ accessToken, selectedTags, useEffect(() => { fetchPerUserData(); - }, [accessToken, selectedTags, currentPage]); - - const handleNextPage = () => { - if (currentPage < perUserData.total_pages) { - setCurrentPage(currentPage + 1); - } - }; - - const handlePrevPage = () => { - if (currentPage > 1) { - setCurrentPage(currentPage - 1); - } - }; + }, [accessToken, selectedTags, pagination.pageIndex, pagination.pageSize]); const columns: ColumnDef[] = [ { @@ -137,30 +127,15 @@ const PerUserUsage: React.FC = ({ accessToken, selectedTags, row.user_id} + paginationMode="server" + pagination={pagination} + onPaginationChange={setPagination} + rowCount={perUserData.total_count} noDataMessage="No per-user usage data" size="compact" /> - - {perUserData.results.length > 10 && ( -
-

Showing 10 of {perUserData.total_count} results

-
- - -
-
- )}
{/* Tab 2: Usage Distribution Histogram */} diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index f6364b5d9d1..546c3b4e018 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -587,6 +587,7 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded model.model_group || String(index)} sortingMode="client" @@ -656,6 +657,7 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded agent.name || String(index)} sortingMode="client" @@ -722,6 +724,7 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded server.server_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx index fce887fc63b..1d2ef75361f 100644 --- a/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx +++ b/ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx @@ -64,6 +64,7 @@ const RoutingGroupsTable: React.FC = ({ return ( group.group_name} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/team/AvailableTeamsTable.tsx b/ui/litellm-dashboard/src/components/team/AvailableTeamsTable.tsx index 6719cc09780..11c430d05d8 100644 --- a/ui/litellm-dashboard/src/components/team/AvailableTeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/team/AvailableTeamsTable.tsx @@ -46,6 +46,7 @@ const AvailableTeamsTable: React.FC = ({ teams, isLoad return ( team.team_id || String(index)} sortingMode="client" diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index be0d0049c13..b10b2584548 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -148,13 +148,60 @@ describe("RequestLogsPanel", () => { }); describe("server-grouped session pagination (#38060)", () => { - it("requests session-grouped pages of 10 rows by default without a cursor", async () => { + it("requests session-grouped pages of 25 rows by default without a cursor", async () => { renderPanel(); await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); expect(lastCall()?.params?.group_by_session).toBe(true); expect(lastCall()?.params?.session_cursor).toBeUndefined(); - expect(lastCall()?.page_size).toBe(10); + expect(lastCall()?.page_size).toBe(25); + }); + + it("offers the same page sizes as the other tables", async () => { + const user = userEvent.setup(); + respondWith([logEntry({ request_id: "req-a" })]); + renderPanel(); + + await waitFor(() => expect(row("req-a")).not.toBeNull()); + await user.click(screen.getByTestId("pagination-page-size")); + + const options = await screen.findAllByRole("option"); + expect(options.map((option) => option.textContent)).toEqual(["25", "50", "100"]); + }); + + it("counts the rendered rows in the footer instead of the server's session total", async () => { + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: [logEntry({ request_id: "req-a" }), logEntry({ request_id: "req-b" }), logEntry({ request_id: "req-c" })], + total: 40, + page: 1, + page_size: 25, + total_pages: 2, + next_session_cursor: null, + has_more: false, + }); + renderPanel(); + + await waitFor(() => expect(row("req-a")).not.toBeNull()); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-3 of 3"); + expect(screen.getByTestId("pagination-next")).toBeDisabled(); + }); + + it("keeps Next enabled from the server total while more session pages remain", async () => { + const firstPage = Array.from({ length: 25 }, (_, index) => logEntry({ request_id: `req-${index}` })); + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: firstPage, + total: 80, + page: 1, + page_size: 25, + total_pages: 4, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 80"); + expect(screen.getByTestId("pagination-next")).toBeEnabled(); }); it("renders every row the server returns without client-side collapsing", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 9b99c6af923..6e984297bf2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -6,13 +6,13 @@ import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } fr import moment from "moment"; import { useCallback, useEffect, useMemo, useState } from "react"; +import { DEFAULT_PAGE_SIZE_OPTIONS } from "@/components/shared/DataTable"; import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import type { KeyResponse } from "../key_team_helpers/key_list"; import { keyInfoV1Call, uiSpendLogsCall } from "../networking"; import KeyInfoView from "../templates/key_info_view"; import type { LogEntry } from "./columns"; -import { LOGS_PAGE_SIZE_OPTIONS } from "./constants"; import { DEFAULT_LOGS_SORTING, formatLogsWindow, @@ -26,7 +26,7 @@ import { LogDetailsDrawer } from "./LogDetailsDrawer"; import { LiveTailBanner, LogsTableToolbar } from "./LogsTableToolbar"; import { RequestLogsTable } from "./RequestLogsTable"; -const PAGE_SIZE = LOGS_PAGE_SIZE_OPTIONS[0]; +const PAGE_SIZE = DEFAULT_PAGE_SIZE_OPTIONS[0]; const DEFAULT_INTERVAL = { value: 24, unit: "hours" }; interface RequestLogsPanelProps { @@ -166,6 +166,10 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const isDrawerOpen = displayLog !== null || displaySessionId !== null; const rows: LogEntry[] = filteredLogs.data; + const rowsThroughThisPage = pagination.pageIndex * pagination.pageSize + rows.length; + const isLastPage = + filteredLogs.has_more === false || (filteredLogs.has_more === undefined && rows.length < pagination.pageSize); + const rowCount = isLastPage ? rowsThroughThisPage : Math.max(filteredLogs.total, rowsThroughThisPage); const handleSearchChange = useCallback((value: string) => { setColumnFilters((previous) => { @@ -290,7 +294,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, Date: Fri, 4 Sep 2026 00:11:00 +0000 Subject: [PATCH 03/61] test(ui): hoist mock responses to named variables to stay within lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/per_user_usage.test.tsx | 3 ++- .../components/view_logs/RequestLogsPanel.test.tsx | 13 +++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/per_user_usage.test.tsx b/ui/litellm-dashboard/src/components/per_user_usage.test.tsx index 5cd0591bb15..01494ef8bfa 100644 --- a/ui/litellm-dashboard/src/components/per_user_usage.test.tsx +++ b/ui/litellm-dashboard/src/components/per_user_usage.test.tsx @@ -80,7 +80,8 @@ describe("PerUserUsage", () => { it("shows every fetched row with a footer that matches the server total and page size", async () => { const results = Array.from({ length: 25 }, (_, index) => userRow(`user-${index}`, "curl/8.0", index)); - mockPerUserAnalyticsCall.mockResolvedValue({ ...mockResponse, results, total_count: 60, total_pages: 3 }); + const firstPage = { ...mockResponse, results, total_count: 60, total_pages: 3 }; + mockPerUserAnalyticsCall.mockResolvedValue(firstPage); render(); await waitFor(() => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index b10b2584548..295446186b1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -170,7 +170,7 @@ describe("RequestLogsPanel", () => { }); it("counts the rendered rows in the footer instead of the server's session total", async () => { - vi.mocked(uiSpendLogsCall).mockResolvedValue({ + const lastPage = { data: [logEntry({ request_id: "req-a" }), logEntry({ request_id: "req-b" }), logEntry({ request_id: "req-c" })], total: 40, page: 1, @@ -178,7 +178,8 @@ describe("RequestLogsPanel", () => { total_pages: 2, next_session_cursor: null, has_more: false, - }); + }; + vi.mocked(uiSpendLogsCall).mockResolvedValue(lastPage); renderPanel(); await waitFor(() => expect(row("req-a")).not.toBeNull()); @@ -187,16 +188,16 @@ describe("RequestLogsPanel", () => { }); it("keeps Next enabled from the server total while more session pages remain", async () => { - const firstPage = Array.from({ length: 25 }, (_, index) => logEntry({ request_id: `req-${index}` })); - vi.mocked(uiSpendLogsCall).mockResolvedValue({ - data: firstPage, + const firstPage = { + data: Array.from({ length: 25 }, (_, index) => logEntry({ request_id: `req-${index}` })), total: 80, page: 1, page_size: 25, total_pages: 4, next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", has_more: true, - }); + }; + vi.mocked(uiSpendLogsCall).mockResolvedValue(firstPage); renderPanel(); await waitFor(() => expect(row("req-0")).not.toBeNull()); From c911740d8292b12c9c7ad4cca926b513127cbc86 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 4 Sep 2026 00:19:13 +0000 Subject: [PATCH 04/61] test(ui): update useModelsInfo call assertions for the new filter arguments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/app/(dashboard)/hooks/models/useModels.test.ts | 4 ++++ 1 file changed, 4 insertions(+) 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 7231c126a63..cfafe82ee30 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 @@ -119,6 +119,8 @@ describe("useModelsInfo", () => { // every other consumer of this hook keeps seeing auto-routers. false, undefined, + undefined, + false, ); expect(modelInfoCall).toHaveBeenCalledTimes(1); }); @@ -147,6 +149,8 @@ describe("useModelsInfo", () => { // every other consumer of this hook keeps seeing auto-routers. false, undefined, + undefined, + false, ); }); From 9e0659212a02dccfdaf74711bffb75bef2a4fda0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 4 Sep 2026 13:48:42 -0700 Subject: [PATCH 05/61] test(e2e): repair two suites broken by intentional behaviour changes Both of these are e2e assumptions that PRs #31731 and #39532 invalidated, not product regressions. They have been red in litellm-e2e builds 119-123. Wildcard readiness probe (6 errors in test_model_access_group_e2e.py) #31731 made _get_wildcard_models drop a wildcard route from /v1/models unconditionally; before it, a wildcard with a matching router deployment stayed in the list and only the no-router / no-deployment fallbacks removed it. The shared readiness helper polls /v1/models for an exact id match, so registering openai/gpt-5.4* now times out at model_servable_timeout every run and every test in the class errors in setup. return_wildcard_routes=True still re-adds the route, so the poll asks for it. The flag is a no-op for a concrete model name -- it only ever adds wildcard entries -- so it is set unconditionally rather than sniffing the name. Semantic auto-router spend assertion #39532 bills the routing embedding to the caller's key on purpose, so the key's spend logs now legitimately carry an openai/text-embedding-3-small row and _assert_served_only_by rejects it. Widening the allowlist would have weakened the assertion this test exists for -- that the request reached the target deployment. Instead the embedding row is split off and asserted separately, which turns the break into coverage for #39532. The poll gains a predicate so it waits for the embedding row rather than racing whichever row is written first. --- tests/e2e/models.py | 9 +++++++++ tests/e2e/proxy_client.py | 3 ++- .../router/test_auto_router_regressions_e2e.py | 15 +++++++++++++-- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 79d9e011f7e..b5229744d6f 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -851,6 +851,15 @@ class ModelListEntry(BaseModel): id: str +class ModelsListParams(BaseModel): + """Query for GET /v1/models. A wildcard route such as ``openai/gpt-5.4*`` is + listed only under ``return_wildcard_routes``; without it the route is dropped + and only its expansions remain, so a readiness poll for the pattern itself + never resolves.""" + + return_wildcard_routes: bool = True + + class ModelsListResponse(BaseModel): """GET /v1/models on the data plane: the deployments the gateway can actually serve right now. Used to confirm a freshly created model has propagated from diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 2d382a610e1..cdc20e5299a 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -55,6 +55,7 @@ from models import ( ModelMode, ModelNewBody, ModelNewResponse, + ModelsListParams, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -336,7 +337,7 @@ class ProxyClient: lambda poll_timeout: self.transport.get( "/v1/models", headers=headers, - params=NoBody(), + params=ModelsListParams(), response_type=ModelsListResponse, timeout=poll_timeout, ), diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 35ba2c8d3d1..188db2a8eb5 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -597,9 +597,20 @@ class TestSemanticAutoRouterResponses: ) ) assert answer.id, "/v1/responses through the semantic auto-router returned no response id" - rows: Final = proxy.poll_logs_for_key(key, min_rows=1) + rows: Final = proxy.poll_logs_for_key( + key, + min_rows=2, + predicate=lambda logged: any(row.model == EMBEDDING_MODEL for row in logged), + ) + embedding_rows: Final = tuple(row for row in rows if row.model == EMBEDDING_MODEL) + assert embedding_rows, ( + "the routing embedding was not billed to the caller's key; " + f"spend logs show {tuple(row.model for row in rows)}" + ) _assert_served_only_by( - rows, CHEAP_SERVED | {semantic_auto_router.target}, "semantic auto-router /v1/responses string input" + [row for row in rows if row.model != EMBEDDING_MODEL], + CHEAP_SERVED | {semantic_auto_router.target}, + "semantic auto-router /v1/responses string input", ) From 98a0cf306f213f511744502b22ed3f3a2a00d5bc Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 15:17:00 -0700 Subject: [PATCH 06/61] fix(shadow_eval): size the judge output cap for a judge that reasons The cap covers reasoning tokens as well as the verdict, and the models people pick as judges reason before answering whether the call asks them to or not: Anthropic's 5 family thinks adaptively and cannot be told not to, so the reasoning bills against max_tokens with nothing in the request to opt out. At 1500 the reasoning consumed the budget and the reply arrived empty or cut off mid-object, which the attempt recorded as an unparseable judge verdict rather than a result. Headroom costs nothing: max_tokens is a ceiling and only generated tokens bill, so the only movement is that judge calls which used to bill their full budget and return nothing now return a verdict. Deliberately not passing reasoning_effort to bound the reasoning instead: is_thinking_enabled treats any reasoning_effort as thinking-enabled, which drops the forced tool_choice that json_mode relies on and turns thinking on with a 1024-token floor for judges that were not reasoning at all. --- litellm/integrations/shadow_eval_logger.py | 10 ++-- .../integrations/test_shadow_eval_logger.py | 49 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 27da785331a..a1716c0954d 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,9 +60,13 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object; a tighter budget truncates the JSON -# mid-object and the attempt is lost to an error row. -JUDGE_MAX_OUTPUT_TOKENS: Final = 1500 +# The judge answers with a small JSON object, but the cap covers reasoning tokens too, +# and the models people pick as judges reason before answering whether or not the call +# asks them to (Anthropic's 5 family thinks adaptively and cannot be told not to). A +# budget sized for the JSON alone is spent on invisible reasoning instead, and the reply +# arrives empty or truncated mid-object, which the attempt records as an unparseable +# verdict. Headroom is free: max_tokens is a ceiling, and only generated tokens bill. +JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 5628d69de26..877677505d6 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -120,6 +120,27 @@ def _router( return router +def _reasoning_judge_router(reasoning_tokens, verdict='{"preference": "A", "confidence": 0.9}'): + """A router whose judge arm reasons before it answers, the way Anthropic's 5 family + does whether or not the call asks it to. Reasoning is billed against the caller's own + max_tokens and the reply is cut off at that cap, so a cap that does not clear the + reasoning budget yields a truncated verdict or no verdict at all. One character stands + in for one token, which is what makes the cap the thing under test.""" + router = MagicMock() + router.model_group_alias = {} + router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + async def acompletion(**kwargs): + if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: + kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} + return {"choices": [{"message": {"content": "shadow answer"}}]} + budget_for_the_answer = kwargs["max_tokens"] - reasoning_tokens + return {"choices": [{"message": {"content": verdict[: max(0, budget_for_the_answer)]}}]} + + router.acompletion = MagicMock(side_effect=acompletion) + return router + + def _spend_counter(store=None): """In-memory stand-in for the proxy's cross-pod spend counter: reads take the max of the counter and the caller's fallback, exactly like get_current_spend does for a key @@ -1134,6 +1155,34 @@ class TestShadowPipeline: assert row["shadow_cost"] == 0.007 assert logger._test_counter["spend:shadow_eval:job-1"] == 0.007 + async def test_the_judge_output_cap_leaves_room_for_a_reasoning_judge(self): + """The output cap covers reasoning tokens as well as the answer, and the models + people pick as judges reason before answering whether or not the call asks them to. + A cap sized for the verdict JSON alone is spent on reasoning instead and the reply + arrives empty, which the attempt records as an unparseable verdict rather than a + result. The judge here burns a reasoning budget typical of a thinking model on a + comparison task, so the cap has to clear it for the verdict to survive.""" + reasoning_tokens = 2000 + logger = _logger(router=_reasoning_judge_router(reasoning_tokens), prisma=(prisma := _prisma())) + + await logger._run_shadow_eval( + job=_job(), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + + row = prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"] + assert row["outcome"] in ("real", "shadow", "tie"), row["error"] + assert row["error"] is None + async def test_a_pipeline_error_after_the_shadow_call_keeps_its_billed_cost(self, monkeypatch: pytest.MonkeyPatch): """An unexpected error between the billed shadow call and the attempt write must still record the shadow cost, or the per-key dollar gate undercounts forever.""" From a2f926eb8f7fac36a193a851b035eabb92b27373 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 15:55:19 -0700 Subject: [PATCH 07/61] fix(shadow_eval): correct the judge output cap's causal claim The prior commit claimed claude-sonnet-5 reasons invisibly by default and eats the judge's budget regardless of what the call asks for. Verified against a live proxy: with no thinking param (what _call_judge sends today), forced tool-choice json_mode, native structured output, and even an explicit thinking=adaptive, the model returned 0 reasoning tokens and a clean compact verdict every time, on prompts up to several thousand characters. The real mechanism only shows up with an elevated reasoning_effort or output_config.effort on the request, which happens when the judge_model deployment is configured with one, e.g. an admin pointing the judge at their best reasoning model. Reproduced directly: reasoning_effort=max, 300-token cap, real Anthropic reply came back finish_reason=length, content=None, 299 of 300 tokens spent on reasoning. Same request at 4096 returned a valid verdict. This is a narrower, verified claim than the one it replaces. --- litellm/integrations/shadow_eval_logger.py | 12 +++++----- .../integrations/test_shadow_eval_logger.py | 22 +++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index a1716c0954d..b554c4bc668 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,12 +60,12 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object, but the cap covers reasoning tokens too, -# and the models people pick as judges reason before answering whether or not the call -# asks them to (Anthropic's 5 family thinks adaptively and cannot be told not to). A -# budget sized for the JSON alone is spent on invisible reasoning instead, and the reply -# arrives empty or truncated mid-object, which the attempt records as an unparseable -# verdict. Headroom is free: max_tokens is a ceiling, and only generated tokens bill. +# The judge answers with a small JSON object, but the cap covers reasoning tokens too. A +# judge_model deployment configured with an elevated reasoning_effort or thinking budget +# (a realistic pick: an admin's best reasoning model doubling as the judge) spends most or +# all of a tight cap on that reasoning, invisibly to this call, and the reply arrives empty +# or truncated mid-object, which the attempt records as an unparseable verdict. Headroom is +# free: max_tokens is a ceiling, and only generated tokens bill. JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 877677505d6..367ad758772 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -121,11 +121,11 @@ def _router( def _reasoning_judge_router(reasoning_tokens, verdict='{"preference": "A", "confidence": 0.9}'): - """A router whose judge arm reasons before it answers, the way Anthropic's 5 family - does whether or not the call asks it to. Reasoning is billed against the caller's own - max_tokens and the reply is cut off at that cap, so a cap that does not clear the - reasoning budget yields a truncated verdict or no verdict at all. One character stands - in for one token, which is what makes the cap the thing under test.""" + """A router whose judge arm reasons before it answers, the way a deployment carrying an + elevated reasoning_effort does. Reasoning is billed against the caller's own max_tokens + and the reply is cut off at that cap, so a cap that does not clear the reasoning budget + yields a truncated verdict or no verdict at all. One character stands in for one token, + which is what makes the cap the thing under test.""" router = MagicMock() router.model_group_alias = {} router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) @@ -1156,12 +1156,12 @@ class TestShadowPipeline: assert logger._test_counter["spend:shadow_eval:job-1"] == 0.007 async def test_the_judge_output_cap_leaves_room_for_a_reasoning_judge(self): - """The output cap covers reasoning tokens as well as the answer, and the models - people pick as judges reason before answering whether or not the call asks them to. - A cap sized for the verdict JSON alone is spent on reasoning instead and the reply - arrives empty, which the attempt records as an unparseable verdict rather than a - result. The judge here burns a reasoning budget typical of a thinking model on a - comparison task, so the cap has to clear it for the verdict to survive.""" + """The output cap covers reasoning tokens as well as the answer, and a judge_model + deployment carrying an elevated reasoning_effort spends that budget before it writes + anything. A cap sized for the verdict JSON alone goes entirely to reasoning and the + reply arrives empty, which the attempt records as an unparseable verdict rather than + a result. The judge here burns a reasoning budget a live claude-sonnet-5 call was + measured at, so the cap has to clear it for the verdict to survive.""" reasoning_tokens = 2000 logger = _logger(router=_reasoning_judge_router(reasoning_tokens), prisma=(prisma := _prisma())) From dd60b7e40f44d7687ad4577332d6f57fc21d7af3 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:16:38 -0700 Subject: [PATCH 08/61] feat(auto-router): decouple compression between the routing decision and the model call An auto router marker deployment can now set auto_router_routing_compression and auto_router_model_compression in its litellm_params, naming the compression guardrail each hop should use (or "none" for no compression on that hop). Neither key set means the request's own compression guardrails keep applying to both hops unchanged. Backend: Router.async_pre_routing_hook resolves the marker's policy and compresses a copy of the messages for the routing decision only when the policy differs from what the model call already got; when both hops share the same compression, it reuses what the ordinary pre-call guardrail pipeline already produced instead of compressing twice. The proxy layer suppresses every other compression guardrail once a policy is engaged and arms the model-side guardrail even when it is not default_on. UI: the auto router's Detailed Configuration gains an Advanced: Compression section with a routing-decision selector and a same/different toggle for the model call, matching the same/different address pattern. --- litellm/constants.py | 4 + litellm/integrations/custom_guardrail.py | 18 ++ litellm/proxy/common_request_processing.py | 7 + .../guardrails/auto_router_compression.py | 211 +++++++++++++ litellm/router.py | 53 +++- litellm/types/router.py | 4 + litellm/types/utils.py | 2 + .../integrations/test_custom_guardrail.py | 44 +++ .../test_auto_router_compression.py | 285 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 54 ++++ tests/test_litellm/test_router.py | 151 ++++++++++ .../add_model/ComplexityRouterConfig.tsx | 34 +++ .../add_model/CompressionControls.tsx | 93 ++++++ .../add_model/add_auto_router_tab.test.tsx | 75 ++++- .../add_model/add_auto_router_tab.tsx | 11 + .../buildAutoRouterCompression.test.ts | 93 ++++++ .../add_model/buildAutoRouterCompression.ts | 52 ++++ .../handle_add_auto_router_submit.tsx | 5 +- .../edit_auto_router_modal.test.tsx | 81 +++++ .../edit_auto_router_modal.tsx | 18 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 + 21 files changed, 1297 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/guardrails/auto_router_compression.py create mode 100644 tests/test_litellm/proxy/guardrails/test_auto_router_compression.py create mode 100644 ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts create mode 100644 ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts diff --git a/litellm/constants.py b/litellm/constants.py index da731cb5eb2..25fdaec20de 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -215,6 +215,10 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # so the deployment-level hook does not re-run them for the same request PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" +# Metadata key listing compression guardrails an auto router's own compression +# policy suppresses for this request. See litellm.proxy.guardrails.auto_router_compression. +AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: Final = "_auto_router_suppressed_compression_guardrails" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9bb613654e4..7f8effa2317 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -45,6 +45,7 @@ dc: Final = DualCache() from litellm.constants import ( + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY, GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) @@ -940,6 +941,20 @@ class CustomGuardrail(CustomLogger): """ return False + def _suppressed_by_auto_router_compression(self, data: dict) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. + + Set only by litellm.proxy.guardrails.auto_router_compression.arm_pre_call, never + by the caller, so a request cannot suppress its own guardrails this way. + """ + for meta_key in ("metadata", "litellm_metadata"): + meta = data.get(meta_key) + if isinstance(meta, dict): + suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) + if isinstance(suppressed, list) and self.guardrail_name in suppressed: + return True + return False + def should_run_guardrail( self, data, @@ -948,6 +963,9 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ + if self._suppressed_by_auto_router_compression(data): + return False + requested_guardrails: Final = self.get_guardrail_from_metadata(data) disable_global_guardrail: Final = self.get_disable_global_guardrail(data) opted_out_global_guardrails: Final = self.get_opted_out_global_guardrails_from_metadata(data) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f25fa46197e..534b2db3e61 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails from litellm.router import Router @@ -2004,6 +2005,12 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + # An auto router with its own compression policy is authoritative for this + # request: suppress every other compression guardrail and arm whichever one + # the policy names for the model call, before those guardrails get a chance + # to run below. + self.data = await _arm_auto_router_compression(data=self.data, llm_router=llm_router) + self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py new file mode 100644 index 00000000000..c3fba937d22 --- /dev/null +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -0,0 +1,211 @@ +""" +Decouples prompt compression between an auto router's routing decision and the +model it routes to. An auto router marker deployment may set +``auto_router_routing_compression`` and/or ``auto_router_model_compression`` in its +``litellm_params`` to name the compression guardrail that hop should use, or +``"none"`` to run no compression on that hop. Neither key set means the request's +own compression guardrails (key/team/model-level, or an "Always on" guardrail) +apply to both hops unchanged, exactly as before this feature existed. + +Once either key is set, this auto router is authoritative: every other compression +guardrail is suppressed for that request, and only these two settings decide what +each hop sees. +""" + +import copy +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +from litellm._logging import verbose_proxy_logger +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, +) +from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.router import Router +else: + CustomGuardrail = Any + Router = Any + +COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) +_NO_COMPRESSION: Final = "none" + +# Metadata key stashing the pre-compression messages so a routing decision that +# names a different compression than the model call still compresses the +# original text, not whatever the model-side guardrail already rewrote it to. +AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: Final = "_auto_router_routing_messages_snapshot" + + +@dataclass(frozen=True, slots=True) +class AutoRouterCompressionPolicy: + """An auto router's compression choice for each hop. ``None`` means no compression.""" + + routing: str | None + model: str | None + + @property + def is_same(self) -> bool: + return self.routing == self.model + + +def _normalized_compression_choice(raw: object) -> str | None: + if not isinstance(raw, str) or not raw: + return None + return None if raw.strip().lower() == _NO_COMPRESSION else raw + + +def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRouterCompressionPolicy | None: + raw_routing: Final = litellm_params.get("auto_router_routing_compression") + raw_model: Final = litellm_params.get("auto_router_model_compression") + if raw_routing is None and raw_model is None: + return None + return AutoRouterCompressionPolicy( + routing=_normalized_compression_choice(raw_routing), + model=_normalized_compression_choice(raw_model), + ) + + +def policy_for_model( + llm_router: "Router | None", model_alias: str, team_id: str | None +) -> AutoRouterCompressionPolicy | None: + """The compression policy declared by the auto router marker deployment `model_alias` resolves to. + + Mirrors the alias lookup in ``_check_and_merge_model_level_guardrails``: this runs + before routing has picked a strategy, so it takes the first marker deployment for + the alias rather than disambiguating by request tags. + """ + if llm_router is None: + return None + deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] + for deployment in deployments: + litellm_params: Final = deployment.get("litellm_params") or {} + model_field = litellm_params.get("model") + if not isinstance(model_field, str) or not model_field.startswith(AUTO_ROUTER_MODEL_PREFIX): + continue + policy = policy_from_litellm_params(litellm_params) + if policy is not None: + return policy + return None + + +def _active_compression_guardrail_names() -> frozenset[str]: + """Names of every currently-active guardrail whose type is a compression guardrail.""" + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + compression_classes: Final = tuple( + cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS + ) + if not compression_classes: + return frozenset() + active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + return frozenset( + cb.guardrail_name for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name + ) + + +async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: + """Apply an auto router's compression policy, if any, before guardrails run. + + Suppresses every other compression guardrail, re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``, and + snapshots the pre-compression messages so the routing decision can compress + them independently of whatever the model-side guardrail does to `data`. + """ + if llm_router is None: + return data + + model_alias: Final = data.get("model") + if not isinstance(model_alias, str) or not model_alias: + return data + + # Read-only until a policy is confirmed: creating the metadata bucket for every + # request, including the vast majority with no auto-router compression policy, + # would be an unwanted side effect of merely checking for one. + metadata_key: Final = get_metadata_variable_name_from_kwargs(data) + existing_bucket: Final = data.get(metadata_key) + other_bucket: Final = data.get("metadata" if metadata_key == "litellm_metadata" else "litellm_metadata") + team_id: Final = (existing_bucket.get("user_api_key_team_id") if isinstance(existing_bucket, dict) else None) or ( + other_bucket.get("user_api_key_team_id") if isinstance(other_bucket, dict) else None + ) + + policy: Final = policy_for_model(llm_router=llm_router, model_alias=model_alias, team_id=team_id) + if policy is None: + return data + + _, metadata = get_or_create_metadata_bucket(data) + suppressed: Final = _active_compression_guardrail_names() - ({policy.model} if policy.model else set()) + if suppressed: + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = sorted(suppressed) + + if policy.model is not None: + requested = metadata.get("guardrails") + if isinstance(requested, list): + if policy.model not in requested: + requested.append(policy.model) + else: + metadata["guardrails"] = [policy.model] + + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) + if snapshot is not None: + metadata[AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] = copy.deepcopy(snapshot) + + return data + + +async def messages_for_routing( + policy: AutoRouterCompressionPolicy | None, + messages: list[dict[str, Any]] | None, + request_kwargs: Mapping[str, object], +) -> list[dict[str, Any]] | None: + """Messages to use for a routing decision, compressed per `policy.routing`. + + Returns None when there is no policy or the policy's routing side names no + compression, meaning the caller should route on whatever messages it already + has. The model call is untouched by this function either way: model-side + compression, if any, already ran as an ordinary pre-call guardrail before the + router was ever reached. + """ + if policy is None or policy.routing is None: + return None + + from litellm.proxy.common_utils.registry_read_through import ( + get_initialized_guardrail_with_read_through, + ) + + metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) + metadata: Final = request_kwargs.get(metadata_key) + snapshot: Final = metadata.get(AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY) if isinstance(metadata, dict) else None + original: Final = snapshot if isinstance(snapshot, list) else messages + if not original: + return None + + guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) + if guardrail is None: + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing + ) + return None + + inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} + # A throwaway request_data: apply_guardrail writes its stats onto this dict, not + # the real request's metadata, so routing-side compression never double-counts + # against extract_compression_saved_tokens's model-savings accounting. + throwaway_request_data: Final[dict[str, object]] = { + "messages": original, + "model": request_kwargs.get("model"), + } + result: Final = await guardrail.apply_guardrail( + inputs=inputs, request_data=throwaway_request_data, input_type="request" + ) + compressed = result.get("structured_messages") + return compressed if isinstance(compressed, list) else original diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..8149ec60ddc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13037,13 +13037,46 @@ class Router: ) return None + from litellm.proxy.guardrails.auto_router_compression import ( + messages_for_routing, + policy_from_litellm_params, + ) + + marker_params: Final = self._alias_marker_litellm_params(registered_model_name, selected_strategy.tags) + compression_policy: Final = policy_from_litellm_params(marker_params) if marker_params else None + # When both hops share the same compression, the model-side guardrail already + # ran in the proxy's ordinary pre-call hook and compressed `messages` in place + # (arm_pre_call armed it whether or not it is `default_on`); reuse that result + # for routing too instead of paying for a second compression call against the + # same content. + needs_independent_routing_compression: Final = compression_policy is not None and not ( + compression_policy.is_same and compression_policy.model is not None + ) + routing_messages: Final = ( + await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs) + if needs_independent_routing_compression + else None + ) + pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( model=registered_model_name, request_kwargs=request_kwargs, - messages=messages, + messages=routing_messages if routing_messages is not None else messages, input=input, specific_deployment=specific_deployment, ) + # The strategy only echoes back whatever `messages` it was handed, so a + # routing-only compression must not leak into the response: the model call + # and downstream deployment-context filtering both key off this field. + # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, + # and pydantic reconstructs a validated list field rather than keeping the + # exact object passed in, even when nothing about it changed. + if ( + pre_routing_hook_response is not None + and routing_messages is not None + and pre_routing_hook_response.messages == routing_messages + ): + pre_routing_hook_response = pre_routing_hook_response.model_copy(update={"messages": messages}) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), @@ -13100,9 +13133,16 @@ class Router: return pre_routing_hook_response - def _forwardable_alias_marker_params( + def _alias_marker_litellm_params( self, model: str, strategy_tags: tuple[str, ...] - ) -> tuple[tuple[str, object], ...]: + ) -> Mapping[str, object] | None: + """The auto-router marker deployment's own `litellm_params` for `model`, tag-scoped. + + Shared by `_forwardable_alias_marker_params` (forwarding api_base/api_key/... + gaps onto the routed deployment) and the auto-router compression policy lookup + (reading `auto_router_routing_compression`/`auto_router_model_compression`), so + both read the same marker row when an alias has more than one, tag-scoped marker. + """ marker_params: Final = tuple( litellm_params for idx in self.model_name_to_deployment_indices.get(model, ()) @@ -13112,7 +13152,12 @@ class Router: tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags ) - selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) + return tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) + + def _forwardable_alias_marker_params( + self, model: str, strategy_tags: tuple[str, ...] + ) -> tuple[tuple[str, object], ...]: + selected: Final = self._alias_marker_litellm_params(model, strategy_tags) if selected is None: return () return tuple( diff --git a/litellm/types/router.py b/litellm/types/router.py index 7ebd50f1328..f5295d6569c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -359,6 +359,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): auto_router_default_model: str | None = None auto_router_embedding_model: str | None = None auto_router_max_input_chars: int | None = None + # Compression policy for the two hops of a routed request. Both unset means the + # request's own compression guardrails apply to both, as they always have. + auto_router_routing_compression: str | None = None + auto_router_model_compression: str | None = None # complexity-router params complexity_router_config: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c1d90694f14..a1fda3d0524 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3713,6 +3713,8 @@ all_litellm_params = ( "auto_router_default_model", "auto_router_embedding_model", "auto_router_max_input_chars", + "auto_router_routing_compression", + "auto_router_model_compression", "complexity_router_config", "complexity_router_default_model", "adaptive_router_config", diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6edc2c9bf77..1fb4299cb56 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -518,6 +518,50 @@ class TestCustomGuardrailShouldRunGuardrail: is True ) + def test_should_run_guardrail_suppressed_by_auto_router_compression(self): + """An auto router's own compression policy can suppress an otherwise-eligible + guardrail, even one that is default_on and explicitly requested.""" + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + data = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["headroom-default"], + }, + } + + assert ( + always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + is False + ) + + def test_should_run_guardrail_suppression_list_does_not_affect_other_names(self): + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + data = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["some-other-guardrail"], + }, + } + + assert ( + always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) + is True + ) + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py new file mode 100644 index 00000000000..b2e83e75768 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -0,0 +1,285 @@ +""" +Unit tests for litellm.proxy.guardrails.auto_router_compression. + +Covers: +- policy_from_litellm_params: absent keys mean no policy; the "none" sentinel + normalizes to explicit no-compression within an active policy; is_same +- policy_for_model: finds the auto-router marker deployment for an alias +- arm_pre_call: no-op without a policy; suppresses active compression guardrails; + arms the model-side guardrail even when it isn't default_on; snapshots messages +- messages_for_routing: no-op without a policy or an unset routing side; compresses + via the named guardrail's apply_guardrail; never writes stats onto the caller's + own request_kwargs (regression for double-counted compression savings) +""" + +from typing import Any + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails.auto_router_compression import ( + AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY, + AutoRouterCompressionPolicy, + arm_pre_call, + messages_for_routing, + policy_for_model, + policy_from_litellm_params, +) +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.types.utils import GenericGuardrailAPIInputs + + +class TestPolicyFromLitellmParams: + def test_neither_key_set_is_no_policy(self): + assert policy_from_litellm_params({}) is None + + def test_routing_only(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "headroom-a"}) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_normalizes_to_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"} + ) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_none_sentinel_is_case_insensitive(self): + policy = policy_from_litellm_params({"auto_router_routing_compression": "NONE"}) + assert policy == AutoRouterCompressionPolicy(routing=None, model=None) + + def test_is_same_true_for_matching_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "x"} + ) + assert policy.is_same is True + + def test_is_same_false_for_different_names(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "x", "auto_router_model_compression": "y"} + ) + assert policy.is_same is False + + def test_is_same_true_when_both_no_compression(self): + policy = policy_from_litellm_params( + {"auto_router_routing_compression": "none", "auto_router_model_compression": "none"} + ) + assert policy.is_same is True + + +class _FakeRouter: + """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + + def __init__(self, deployments: list[dict[str, Any]]): + self._deployments = deployments + + def get_model_list(self, model_name, team_id=None): + return [d for d in self._deployments if d.get("model_name") == model_name] + + +class TestPolicyForModel: + def test_no_router_returns_none(self): + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None) is None + + def test_no_marker_deployment_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + + def test_marker_deployment_without_policy_returns_none(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] + ) + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + + def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + +class _RecordingCompressionGuardrail(CustomGuardrail): + """A guardrail whose apply_guardrail marks every text message as compressed.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.request_data_seen: list[dict] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.request_data_seen.append(request_data) + structured_messages = inputs.get("structured_messages") or [] + compressed = [ + {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages + ] + return {**inputs, "structured_messages": compressed} + + +@pytest.fixture +def registered_guardrail(): + import litellm + + guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + +class TestArmPreCall: + @pytest.mark.asyncio + async def test_no_router_is_noop(self): + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=None) + assert result == data + assert "metadata" not in result + + @pytest.mark.asyncio + async def test_no_policy_does_not_create_metadata_bucket(self): + router = _FakeRouter( + [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + assert "metadata" not in result + assert "litellm_metadata" not in result + + @pytest.mark.asyncio + async def test_policy_suppresses_active_compression_guardrails(self, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem( + guardrail_registry.guardrail_class_registry, "fake-provider", _RecordingCompressionGuardrail + ) + monkeypatch.setattr( + "litellm.proxy.guardrails.auto_router_compression.COMPRESSION_GUARDRAIL_PROVIDERS", + frozenset({"fake-provider"}), + ) + import litellm + + always_on = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + litellm.logging_callback_manager.add_litellm_callback(always_on) + try: + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] + assert "always-on-compression" in suppressed + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + + @pytest.mark.asyncio + async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-b", + }, + } + ] + ) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + result = await arm_pre_call(data=data, llm_router=router) + assert result["metadata"]["guardrails"] == ["headroom-b"] + + @pytest.mark.asyncio + async def test_snapshots_original_messages(self): + router = _FakeRouter( + [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "headroom-a", + "auto_router_model_compression": "none", + }, + } + ] + ) + original_messages = [{"role": "user", "content": "hi"}] + data = {"model": "smart-router", "messages": original_messages} + result = await arm_pre_call(data=data, llm_router=router) + snapshot = result["metadata"][AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] + assert snapshot == original_messages + assert snapshot is not original_messages # a copy, not the live reference + + +class TestMessagesForRouting: + @pytest.mark.asyncio + async def test_no_policy_returns_none(self): + assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_routing_side_unset_returns_none(self): + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_returns_none(self): + policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None) + messages = [{"role": "user", "content": "hi"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result is None + + @pytest.mark.asyncio + async def test_compresses_via_the_named_guardrail(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hello world"}] + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}] + + @pytest.mark.asyncio + async def test_uses_the_snapshot_when_present(self, registered_guardrail): + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + snapshot = [{"role": "user", "content": "original"}] + request_kwargs = {"metadata": {AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: snapshot}} + # `messages` here stands in for whatever a model-side guardrail already + # rewrote `data["messages"]` to -- routing must ignore it and compress the + # pristine snapshot instead. + already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] + result = await messages_for_routing( + policy=policy, messages=already_rewritten, request_kwargs=request_kwargs + ) + assert result == [{"role": "user", "content": "[COMPRESSED] original"}] + + @pytest.mark.asyncio + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs( + self, registered_guardrail + ): + """Regression: a real compression guardrail writes its stats onto whatever + `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ + request_data`). If that were the caller's own `request_kwargs`, routing-side + compression would double-count into extract_compression_saved_tokens, which + sums every guardrail_information entry on the real request's metadata.""" + policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) + messages = [{"role": "user", "content": "hi"}] + request_kwargs = {"metadata": {}} + await messages_for_routing(policy=policy, messages=messages, request_kwargs=request_kwargs) + assert registered_guardrail.request_data_seen[0] is not request_kwargs + assert request_kwargs == {"metadata": {}} diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f7fe6ad9d39..c0809e53d2e 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -376,6 +376,60 @@ class TestProxyBaseLLMRequestProcessing: assert "litellm_logging_obj" not in persisted_body json.dumps(persisted_body) + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): + """arm_pre_call must run before pre_call_hook: an auto router's own compression + policy has to be in `data["metadata"]` (naming the model-side guardrail so it + runs even if it isn't default_on) by the time guardrails see the request.""" + processing_obj = ProxyBaseLLMRequestProcessing(data={}) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + + seen_metadata: dict = {} + + async def mock_pre_call_hook(user_api_key_dict, data, call_type): + seen_metadata.update(data.get("metadata") or {}) + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=mock_pre_call_hook) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + + fake_llm_router = MagicMock() + fake_llm_router.get_model_list.return_value = [ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "auto_router_routing_compression": "none", + "auto_router_model_compression": "headroom-model", + }, + } + ] + mock_proxy_config = MagicMock(spec=ProxyConfig) + mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) + + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=fake_llm_router, + ) + + assert seen_metadata.get("guardrails") == ["headroom-model"] + def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch): mock_set_active_span_tag = MagicMock(return_value=True) import litellm.proxy.dd_span_tagger diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5fc96bcfbb1..cdae3b131ae 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -18,6 +18,7 @@ import pytest import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( @@ -9995,6 +9996,156 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestAutoRouterCompressionDecoupling: + """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` + decouple what the routing decision sees from what the model call sees. The one + assertion that must hold under any mutation: the strategy can be routed on + compressed text while the caller's own `messages` list - the one that would reach + the model - is never touched.""" + + class _RecordingStrategy: + """Echoes back whatever `messages` it was handed, like every real strategy does.""" + + def __init__(self): + self.received_messages: list[dict] | None = None + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + self.received_messages = messages + return PreRoutingHookResponse(model="gemini-flash", messages=messages) + + class _CompressingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.call_count = 0 + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.call_count += 1 + structured_messages = inputs.get("structured_messages") or [] + compressed = [ + {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages + ] + return {**inputs, "structured_messages": compressed} + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + def _router(self, marker_litellm_params: dict) -> tuple[litellm.Router, "_RecordingStrategy"]: + from litellm.types.router import TaggedPreRoutingStrategy + + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + router = litellm.Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + **marker_litellm_params, + }, + }, + { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + }, + ], + ) + for name in ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers"): + setattr(router, name, {}) + strategy = self._RecordingStrategy() + router.complexity_routers = {"smart-router": [TaggedPreRoutingStrategy(tags=(), strategy=strategy)]} + return router, strategy + + @pytest.fixture + def registered_guardrail(self): + guardrail = self._CompressingGuardrail(guardrail_name="fake-compress") + litellm.logging_callback_manager.add_litellm_callback(guardrail) + yield guardrail + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) + + @pytest.mark.asyncio + async def test_routing_side_compression_never_reaches_the_caller_messages(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "none", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + assert response.messages == original_messages + + @pytest.mark.asyncio + async def test_model_side_compression_alone_leaves_routing_uncompressed(self, registered_guardrail): + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages == original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): + """The same/different distinction exists so a shared choice does not pay for + compression twice: by the time the router runs, `messages` already reflects + whatever the ordinary pre-call guardrail pipeline did for the model call, so + the routing decision must reuse it rather than calling the guardrail again.""" + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + # Stands in for what the proxy's ordinary pre-call guardrail pipeline would + # have already produced for the model call, since `auto_router_model_compression` + # names a guardrail: the router never triggers that pipeline itself. + already_compressed_messages = [ + {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} + ] + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages + ) + + assert strategy.received_messages == already_compressed_messages + assert response.messages == already_compressed_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.asyncio + async def test_no_policy_is_fully_unaffected(self, registered_guardrail): + router, strategy = self._router({}) + original_messages = self._messages() + + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=original_messages + ) + + assert strategy.received_messages is original_messages + assert response.messages == original_messages + assert registered_guardrail.call_count == 0 + + @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 2a024ab7fdf..42115265034 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -48,6 +48,8 @@ import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; +import CompressionControls from "./CompressionControls"; +import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; export type { DimensionWeights, TierBoundaries, TokenThresholds }; export type { CustomTierSet, TierRow } from "./tier_rows"; @@ -490,6 +492,10 @@ interface ComplexityRouterConfigProps { onMatchThresholdChange?: (threshold: number) => void; escalationKeywords?: string[]; onEscalationKeywordsChange?: (keywords: string[]) => void; + // Optional: not part of complexity_router_config, since it applies to every + // pre-routing strategy, not just the complexity router. + autoRouterCompression?: AutoRouterCompressionState; + onAutoRouterCompressionChange?: (state: AutoRouterCompressionState) => void; showValidationErrors?: boolean; } @@ -611,6 +617,8 @@ const ComplexityRouterConfig: React.FC = ({ onMatchThresholdChange = () => {}, escalationKeywords = [], onEscalationKeywordsChange, + autoRouterCompression = DEFAULT_AUTO_ROUTER_COMPRESSION, + onAutoRouterCompressionChange, showValidationErrors = false, }) => { const customTierSet = value.custom_tier_set; @@ -875,6 +883,32 @@ const ComplexityRouterConfig: React.FC = ({ }, ] : []), + ...(onAutoRouterCompressionChange + ? [ + { + key: "compression", + label: Advanced: Compression, + children: ( + + onAutoRouterCompressionChange({ + ...autoRouterCompression, + routing, + sameAsRouting: routing === undefined ? true : autoRouterCompression.sameAsRouting, + }) + } + sameAsRouting={autoRouterCompression.sameAsRouting} + onSameAsRoutingChange={(sameAsRouting) => + onAutoRouterCompressionChange({ ...autoRouterCompression, sameAsRouting }) + } + model={autoRouterCompression.model} + onModelChange={(model) => onAutoRouterCompressionChange({ ...autoRouterCompression, model })} + /> + ), + }, + ] + : []), ...(onKeywordTierRulesChange || onSemanticMatchingEnabledChange ? [ { diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx new file mode 100644 index 00000000000..a0a240f76b0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -0,0 +1,93 @@ +import { SimpleTooltip } from "@/components/ui/tooltip"; +import { SearchSelect, SearchSelectOption } from "@/components/shared/SearchSelect"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { Info } from "lucide-react"; +import React from "react"; +import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { COMPRESSION_GUARDRAIL_PROVIDER } from "@/app/(dashboard)/cost-optimization/_components/helpers"; +import { NO_COMPRESSION } from "./buildAutoRouterCompression"; + +interface CompressionControlsProps { + routing: string | undefined; + onRoutingChange: (value: string | undefined) => void; + sameAsRouting: boolean; + onSameAsRoutingChange: (same: boolean) => void; + model: string | undefined; + onModelChange: (value: string | undefined) => void; +} + +const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; + +const CompressionControls: React.FC = ({ + routing, + onRoutingChange, + sameAsRouting, + onSameAsRoutingChange, + model, + onModelChange, +}) => { + const { data } = useGuardrails(); + const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) + .filter((g) => (g.litellm_params?.guardrail ?? "").toString().toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER) + .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); + const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; + + return ( +
+
+
+ Routing decision + + + +
+ onRoutingChange(value === "" ? undefined : value)} + placeholder="Inherit from the request's own compression guardrails" + emptyText="No compression guardrails found" + aria-label="Routing decision compression" + /> +
+ + {routing !== undefined && ( +
+ Model call + onSameAsRoutingChange(value === "same")} + className="w-full" + > +
+ + +
+
+ + {!sameAsRouting && ( +
+ onModelChange(value === "" ? undefined : value)} + placeholder="None (no compression)" + emptyText="No compression guardrails found" + aria-label="Model call compression" + /> +
+ )} +
+ )} +
+ ); +}; + +export default CompressionControls; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d2f6b10c3a6..5605a993ded 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -1,4 +1,12 @@ -import { renderWithProviders, screen, waitFor, within, fireEvent, testQueryClient } from "../../../tests/test-utils"; +import { + renderWithProviders, + screen, + waitFor, + within, + fireEvent, + testQueryClient, + chooseSelectOption, +} from "../../../tests/test-utils"; import userEvent from "@testing-library/user-event"; import { vi } from "vitest"; import AddAutoRouterTab from "./add_auto_router_tab"; @@ -522,6 +530,71 @@ describe("AddAutoRouterTab", () => { ); }); + describe("prompt compression", () => { + it("leaves both compression keys out of the create payload when the section is untouched", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-router"); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted).not.toHaveProperty("auto_router_routing_compression"); + expect(submitted).not.toHaveProperty("auto_router_model_compression"); + }); + + it("mirrors an explicit no-compression routing choice onto the model call by default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-explicit-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + + it("defaults the model call to none when different is chosen but nothing is picked there", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "different-compression-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Compression")); + await chooseSelectOption( + user, + screen.getByRole("combobox", { name: "Routing decision compression" }), + "None (no compression)", + ); + await user.click(screen.getByText("Use a different compression")); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + const submitted = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]; + expect(submitted?.auto_router_routing_compression).toBe("none"); + expect(submitted?.auto_router_model_compression).toBe("none"); + }); + }); + // The scalar floor is the one scorer knob with no group dict behind it, so its wiring into the create // payload is only proven end to end. 0 is the case a truthy check would silently drop. it("carries a reasoning override floor of 0 through to the create payload", async () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index a548c2c6533..decacac6501 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -32,6 +32,11 @@ import ComplexityRouterConfig, { } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { DEFAULT_ESCALATION_KEYWORDS } from "./EscalationKeywords"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, +} from "./buildAutoRouterCompression"; import { DEFAULT_MATCH_THRESHOLD } from "./SemanticKeywordMatching"; import { BuildComplexityRouterConfigParams, @@ -194,6 +199,9 @@ const AddAutoRouterTab: React.FC = ({ const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); const [escalationKeywords, setEscalationKeywords] = useState(DEFAULT_ESCALATION_KEYWORDS); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [showValidationErrors, setShowValidationErrors] = useState(false); const [editingTiers, setEditingTiers] = useState(false); const [isSubmitting, setIsSubmitting] = useState(false); @@ -461,6 +469,7 @@ const AddAutoRouterTab: React.FC = ({ model_type: "complexity_router", complexity_router_config: complexityRouterConfigPayload, model_access_group: form.getValues("model_access_group"), + ...buildAutoRouterCompressionParams(autoRouterCompression), }; await handleAddAutoRouterSubmit(submitValues, accessToken, () => form.reset(EMPTY_FORM_VALUES), handleOk); @@ -666,6 +675,8 @@ const AddAutoRouterTab: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} showValidationErrors={showValidationErrors} />
diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts new file mode 100644 index 00000000000..b917fcedaa2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -0,0 +1,93 @@ +import { + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; + +describe("buildAutoRouterCompressionParams", () => { + it("omits both keys when routing was never configured", () => { + expect(buildAutoRouterCompressionParams(DEFAULT_AUTO_ROUTER_COMPRESSION)).toEqual({}); + }); + + it("mirrors routing onto model when same-as-routing is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + }); + + it("uses the explicit model choice when different is chosen", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: "headroom-b", + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + }); + + it("defaults the model side to none when different is chosen but nothing is picked", () => { + const params = buildAutoRouterCompressionParams({ + routing: "headroom-a", + sameAsRouting: false, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: NO_COMPRESSION, + }); + }); + + it("sends the none sentinel when routing itself is explicitly turned off", () => { + const params = buildAutoRouterCompressionParams({ + routing: NO_COMPRESSION, + sameAsRouting: true, + model: undefined, + }); + expect(params).toEqual({ + auto_router_routing_compression: NO_COMPRESSION, + auto_router_model_compression: NO_COMPRESSION, + }); + }); +}); + +describe("hydrateAutoRouterCompression", () => { + it("returns the default state when neither key is set", () => { + expect(hydrateAutoRouterCompression({})).toEqual(DEFAULT_AUTO_ROUTER_COMPRESSION); + }); + + it("is same-as-routing when the model value matches routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("is different when the model value diverges from routing", () => { + const state = hydrateAutoRouterCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-b", + }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" }); + }); + + it("treats a missing model key as same-as-routing", () => { + const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + }); + + it("round-trips through buildAutoRouterCompressionParams", () => { + const original = { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "none" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(original)); + expect(rebuilt).toEqual(original); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts new file mode 100644 index 00000000000..49180d5b4ce --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -0,0 +1,52 @@ +/** + * Maps the auto router's compression form state to the two flat litellm_params keys + * the backend reads (litellm.proxy.guardrails.auto_router_compression), and back. + * + * `routing` being undefined means the section was never touched: both keys are + * omitted from the payload, and the request's own compression guardrails apply to + * both hops unchanged. Once `routing` has a value (a guardrail name, or the "none" + * sentinel for explicit no-compression), the auto router is authoritative and the + * model side always gets a concrete value too, mirroring `routing` when same-as + * is chosen and defaulting to "none" otherwise. + */ + +export const NO_COMPRESSION = "none"; + +export interface AutoRouterCompressionState { + routing: string | undefined; + sameAsRouting: boolean; + model: string | undefined; +} + +export interface AutoRouterCompressionLitellmParams { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; +} + +export const DEFAULT_AUTO_ROUTER_COMPRESSION: AutoRouterCompressionState = { + routing: undefined, + sameAsRouting: true, + model: undefined, +}; + +export const buildAutoRouterCompressionParams = ( + state: AutoRouterCompressionState, +): AutoRouterCompressionLitellmParams => { + if (state.routing === undefined) return {}; + return { + auto_router_routing_compression: state.routing, + auto_router_model_compression: state.sameAsRouting ? state.routing : (state.model ?? NO_COMPRESSION), + }; +}; + +export const hydrateAutoRouterCompression = (litellmParams: { + auto_router_routing_compression?: string | null; + auto_router_model_compression?: string | null; +}): AutoRouterCompressionState => { + const routing = litellmParams.auto_router_routing_compression ?? undefined; + if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + + const model = litellmParams.auto_router_model_compression ?? undefined; + const sameAsRouting = model === undefined || model === routing; + return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; +}; diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx index 9385836ce1a..59d9ecf205e 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_auto_router_submit.tsx @@ -1,8 +1,9 @@ import { modelCreateCall } from "../networking"; import { toast } from "@/lib/toast"; import type { ComplexityRouterConfigPayload } from "./build_complexity_router_config"; +import type { AutoRouterCompressionLitellmParams } from "./buildAutoRouterCompression"; -export interface AddAutoRouterValues { +export interface AddAutoRouterValues extends AutoRouterCompressionLitellmParams { auto_router_name: string; auto_router_default_model: string | undefined; model_type: "complexity_router"; @@ -24,6 +25,8 @@ export const handleAddAutoRouterSubmit = async ( model: "auto_router/complexity_router", complexity_router_config: values.complexity_router_config, complexity_router_default_model: values.auto_router_default_model, + auto_router_routing_compression: values.auto_router_routing_compression, + auto_router_model_compression: values.auto_router_model_compression, }, model_info: { ...(values.team_id ? { team_id: values.team_id } : {}), diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 970bcaa545f..b93db8d963e 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -1029,3 +1029,84 @@ describe("EditAutoRouterModal with a stored custom tier set", () => { expect(savedConfig().tier_model_configs).toEqual(CUSTOM_STORED.tier_model_configs); }); }); + +describe("EditAutoRouterModal prompt compression", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const savedLitellmParams = () => { + const [, payload] = modelPatchUpdateCall.mock.calls.at(-1) ?? []; + return payload?.litellm_params; + }; + + const renderWithStoredCompression = ( + compression?: { auto_router_routing_compression?: string; auto_router_model_compression?: string }, + ) => + renderWithProviders( + , + ); + + it("leaves both compression keys out of an untouched save when none were stored", async () => { + const user = userEvent.setup(); + renderWithStoredCompression(); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_routing_compression"); + expect(savedLitellmParams()).not.toHaveProperty("auto_router_model_compression"); + }); + + it("preserves a stored same-as-routing compression through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "headroom-a", + }); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); + expect(savedLitellmParams()?.auto_router_model_compression).toBe("headroom-a"); + }); + + it("shows a stored different-compression choice as Use a different compression, not Same", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByText("Advanced: Compression")); + + expect(await screen.findByRole("combobox", { name: "Routing decision compression" })).toHaveValue("headroom-a"); + expect(screen.getByRole("radio", { name: "Use a different compression" })).toBeChecked(); + expect(screen.getByRole("combobox", { name: "Model call compression" })).toHaveValue("None (no compression)"); + }); + + it("preserves a stored different-compression choice through an untouched open-and-save", async () => { + const user = userEvent.setup(); + renderWithStoredCompression({ + auto_router_routing_compression: "headroom-a", + auto_router_model_compression: "none", + }); + + await user.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedLitellmParams()?.auto_router_routing_compression).toBe("headroom-a"); + expect(savedLitellmParams()?.auto_router_model_compression).toBe("none"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index ea1e5cba6a3..a4852f9e784 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -41,6 +41,12 @@ import { } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; import { DEFAULT_MATCH_THRESHOLD } from "../add_model/SemanticKeywordMatching"; +import { + type AutoRouterCompressionState, + buildAutoRouterCompressionParams, + DEFAULT_AUTO_ROUTER_COMPRESSION, + hydrateAutoRouterCompression, +} from "../add_model/buildAutoRouterCompression"; import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords"; import { hydrateDimensionWeights, @@ -424,6 +430,9 @@ const EditAutoRouterModal: React.FC = ({ const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); const [embeddingModel, setEmbeddingModel] = useState(undefined); const [matchThreshold, setMatchThreshold] = useState(DEFAULT_MATCH_THRESHOLD); + const [autoRouterCompression, setAutoRouterCompression] = useState( + DEFAULT_AUTO_ROUTER_COMPRESSION, + ); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", @@ -516,6 +525,12 @@ const EditAutoRouterModal: React.FC = ({ setMatchThreshold( typeof parsedConfig.match_threshold === "number" ? parsedConfig.match_threshold : DEFAULT_MATCH_THRESHOLD, ); + setAutoRouterCompression( + hydrateAutoRouterCompression({ + auto_router_routing_compression: modelData.litellm_params?.auto_router_routing_compression, + auto_router_model_compression: modelData.litellm_params?.auto_router_model_compression, + }), + ); form.reset({ ...EMPTY_FORM_VALUES, @@ -628,6 +643,7 @@ const EditAutoRouterModal: React.FC = ({ ...modelData.litellm_params, complexity_router_config: updatedConfig, complexity_router_default_model: defaultModel, + ...buildAutoRouterCompressionParams(autoRouterCompression), }; const updatedModelInfo = { ...modelData.model_info, @@ -749,6 +765,8 @@ const EditAutoRouterModal: React.FC = ({ onMatchThresholdChange={setMatchThreshold} escalationKeywords={escalationKeywords} onEscalationKeywordsChange={setEscalationKeywords} + autoRouterCompression={autoRouterCompression} + onAutoRouterCompressionChange={setAutoRouterCompression} /> ) : ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5f3cca49644..8f6a0700517 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29222,6 +29222,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */ @@ -39275,6 +39279,10 @@ export interface components { auto_router_embedding_model?: string | null; /** Auto Router Max Input Chars */ auto_router_max_input_chars?: number | null; + /** Auto Router Model Compression */ + auto_router_model_compression?: string | null; + /** Auto Router Routing Compression */ + auto_router_routing_compression?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; /** Aws Batch Role Arn */ From 2f5bfae1a61b0821b6af9eabb045522adfa7b28a Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:21:13 -0700 Subject: [PATCH 09/61] refactor(shadow_eval): tighten the judge cap comment and type the test helper --- litellm/integrations/shadow_eval_logger.py | 9 +++------ .../integrations/test_shadow_eval_logger.py | 10 +++++----- 2 files changed, 8 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index b554c4bc668..2c56ecb8721 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,12 +60,9 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object, but the cap covers reasoning tokens too. A -# judge_model deployment configured with an elevated reasoning_effort or thinking budget -# (a realistic pick: an admin's best reasoning model doubling as the judge) spends most or -# all of a tight cap on that reasoning, invisibly to this call, and the reply arrives empty -# or truncated mid-object, which the attempt records as an unparseable verdict. Headroom is -# free: max_tokens is a ceiling, and only generated tokens bill. +# The judge answers with a small JSON object, but the cap covers reasoning tokens too: a +# judge deployment carrying an elevated reasoning_effort spends a tight cap before it ever +# answers, and the truncated reply is recorded as an unparseable verdict. JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 367ad758772..9fcbd116f63 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -120,12 +120,12 @@ def _router( return router -def _reasoning_judge_router(reasoning_tokens, verdict='{"preference": "A", "confidence": 0.9}'): +def _reasoning_judge_router( + reasoning_tokens: int, verdict: str = '{"preference": "A", "confidence": 0.9}' +) -> MagicMock: """A router whose judge arm reasons before it answers, the way a deployment carrying an - elevated reasoning_effort does. Reasoning is billed against the caller's own max_tokens - and the reply is cut off at that cap, so a cap that does not clear the reasoning budget - yields a truncated verdict or no verdict at all. One character stands in for one token, - which is what makes the cap the thing under test.""" + elevated reasoning_effort does: reasoning bills against the caller's own max_tokens and + the reply is cut off at that cap. One character stands in for one token.""" router = MagicMock() router.model_group_alias = {} router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) From da5e38ce9c1fe2ba4854952eac0cc2a694e1c38b Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:33:16 -0700 Subject: [PATCH 10/61] refactor(shadow_eval): state the cap's constraint without the rationale --- litellm/integrations/shadow_eval_logger.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 2c56ecb8721..fc82ebafe09 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -60,9 +60,8 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16 _MAX_JUDGE_RESPONSE_CHARS: Final = 8_000 _MAX_JUDGE_PROMPT_CHARS: Final = 24_000 -# The judge answers with a small JSON object, but the cap covers reasoning tokens too: a -# judge deployment carrying an elevated reasoning_effort spends a tight cap before it ever -# answers, and the truncated reply is recorded as an unparseable verdict. +# Covers the judge's reasoning tokens as well as its small JSON answer: a judge deployment +# carrying an elevated reasoning_effort spends a tight cap before it ever answers. JUDGE_MAX_OUTPUT_TOKENS: Final = 4096 _MAX_ERROR_CHARS: Final = 500 From 9e286fe94bf18d29b7bb56e3c2f77d114c10acc5 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:42:11 -0700 Subject: [PATCH 11/61] fix(auto-router): close review findings on per-hop compression - Suppression markers now carry the per-process token `_pre_call_marker` already uses, so a caller cannot switch off an always-on PII, content-filter or compression guardrail by naming it in its own request metadata. - Routing set to "none" with the model side compressed now classifies on the pre-compression snapshot instead of the model-side guardrail's output. - Both the proxy's pre-call arming and the router's routing hook resolve the policy through one tag-aware `policy_for_model`, so an alias with several tag-scoped markers can no longer suppress one marker's guardrail and then route under another marker's policy. - The pre-compression snapshot moved from request metadata to a ContextVar: `refresh_proxy_server_request_body_snapshot` copies metadata into `proxy_server_request.body`, which deployments persist, and the snapshot holds the prompt as it was before any masking guardrail rewrote it. - The compression selector lists Compresr guardrails too, not just Headroom. --- litellm/integrations/custom_guardrail.py | 22 ++- .../guardrails/auto_router_compression.py | 148 +++++++++------ litellm/router.py | 32 ++-- .../integrations/test_custom_guardrail.py | 35 +++- .../test_auto_router_compression.py | 177 +++++++++++++----- tests/test_litellm/test_router.py | 29 +++ .../add_model/CompressionControls.tsx | 5 +- .../add_model/buildAutoRouterCompression.ts | 7 + 8 files changed, 322 insertions(+), 133 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7f8effa2317..558e97cfc16 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -941,17 +941,29 @@ class CustomGuardrail(CustomLogger): """ return False - def _suppressed_by_auto_router_compression(self, data: dict) -> bool: - """True when an auto router's own compression policy suppresses this guardrail. + def auto_router_suppression_marker(self) -> str | None: + """The value `arm_pre_call` must write to suppress this guardrail. - Set only by litellm.proxy.guardrails.auto_router_compression.arm_pre_call, never - by the caller, so a request cannot suppress its own guardrails this way. + Carries the per-process token for the same reason `_pre_call_marker` does: a + caller controls request metadata, so a bare guardrail name there would let any + request switch off a PII, content-filter, or compression guardrail for itself. + The token is never sent to the caller, so the marker cannot be forged. """ + name: Final = self.guardrail_name + if not name: + return None + return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" + + def _suppressed_by_auto_router_compression(self, data: dict[str, object]) -> bool: + """True when an auto router's own compression policy suppresses this guardrail.""" + marker: Final = self.auto_router_suppression_marker() + if marker is None: + return False for meta_key in ("metadata", "litellm_metadata"): meta = data.get(meta_key) if isinstance(meta, dict): suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, list) and self.guardrail_name in suppressed: + if isinstance(suppressed, list) and marker in suppressed: return True return False diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index c3fba937d22..7ccd1937543 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -12,34 +12,32 @@ guardrail is suppressed for that request, and only these two settings decide wha each hop sees. """ -import copy -from collections.abc import Mapping +import contextvars +from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY -from litellm.litellm_core_utils.core_helpers import ( - get_metadata_variable_name_from_kwargs, - get_or_create_metadata_bucket, -) +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.router import Router -else: - CustomGuardrail = Any - Router = Any COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# Metadata key stashing the pre-compression messages so a routing decision that -# names a different compression than the model call still compresses the -# original text, not whatever the model-side guardrail already rewrote it to. -AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: Final = "_auto_router_routing_messages_snapshot" +# The pre-compression messages, so a routing decision that does not share the model +# call's compression still classifies on the original text. Deliberately a ContextVar +# rather than a metadata key: `refresh_proxy_server_request_body_snapshot` copies +# metadata into `proxy_server_request.body`, which deployments persist to spend logs, +# and this holds the prompt as it was before any masking guardrail rewrote it. +_routing_messages_snapshot: Final[contextvars.ContextVar[tuple[Mapping[str, object], ...] | None]] = ( + contextvars.ContextVar("litellm_auto_router_routing_messages_snapshot", default=None) +) @dataclass(frozen=True, slots=True) @@ -72,30 +70,51 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( - llm_router: "Router | None", model_alias: str, team_id: str | None + llm_router: "Router | None", + model_alias: str, + team_id: str | None, + request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy declared by the auto router marker deployment `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to. - Mirrors the alias lookup in ``_check_and_merge_model_level_guardrails``: this runs - before routing has picked a strategy, so it takes the first marker deployment for - the alias rather than disambiguating by request tags. + Both the proxy's pre-call arming and the router's routing hook resolve the policy + through here, with the same tag rule, so an alias carrying several tag-scoped + markers can never suppress one marker's guardrail and then route under another + marker's policy. """ if llm_router is None: return None deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] - for deployment in deployments: - litellm_params: Final = deployment.get("litellm_params") or {} - model_field = litellm_params.get("model") - if not isinstance(model_field, str) or not model_field.startswith(AUTO_ROUTER_MODEL_PREFIX): - continue - policy = policy_from_litellm_params(litellm_params) + markers: Final = tuple( + litellm_params + for deployment in deployments + if isinstance(litellm_params := deployment.get("litellm_params") or {}, Mapping) + and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + ) + requested: Final = frozenset(request_tags) + tag_matched: Final = tuple( + params for params in markers if requested.issuperset(frozenset(params.get("tags") or ())) + ) + for params in (*tag_matched, *markers): + policy = policy_from_litellm_params(params) if policy is not None: return policy return None -def _active_compression_guardrail_names() -> frozenset[str]: - """Names of every currently-active guardrail whose type is a compression guardrail.""" +def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + for meta_key in ("metadata", "litellm_metadata"): + meta = request_kwargs.get(meta_key) + if isinstance(meta, Mapping): + team_id = meta.get("user_api_key_team_id") + if isinstance(team_id, str): + return team_id + return None + + +def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: + """Every currently-active guardrail whose type is a compression guardrail.""" import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry @@ -104,21 +123,20 @@ def _active_compression_guardrail_names() -> frozenset[str]: cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS ) if not compression_classes: - return frozenset() + return () active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) - return frozenset( - cb.guardrail_name for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name - ) + return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name) -async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: +async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | None") -> MutableMapping[str, object]: """Apply an auto router's compression policy, if any, before guardrails run. Suppresses every other compression guardrail, re-enables the model-side guardrail the policy names (if any) even when it isn't ``default_on``, and - snapshots the pre-compression messages so the routing decision can compress - them independently of whatever the model-side guardrail does to `data`. + snapshots the pre-compression messages so the routing decision can read them + independently of whatever the model-side guardrail does to `data`. """ + _routing_messages_snapshot.set(None) if llm_router is None: return data @@ -129,21 +147,27 @@ async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: # Read-only until a policy is confirmed: creating the metadata bucket for every # request, including the vast majority with no auto-router compression policy, # would be an unwanted side effect of merely checking for one. - metadata_key: Final = get_metadata_variable_name_from_kwargs(data) - existing_bucket: Final = data.get(metadata_key) - other_bucket: Final = data.get("metadata" if metadata_key == "litellm_metadata" else "litellm_metadata") - team_id: Final = (existing_bucket.get("user_api_key_team_id") if isinstance(existing_bucket, dict) else None) or ( - other_bucket.get("user_api_key_team_id") if isinstance(other_bucket, dict) else None - ) + from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs - policy: Final = policy_for_model(llm_router=llm_router, model_alias=model_alias, team_id=team_id) + policy: Final = policy_for_model( + llm_router=llm_router, + model_alias=model_alias, + team_id=team_id_from_request(data), + request_tags=_get_tags_from_request_kwargs(data), + ) if policy is None: return data _, metadata = get_or_create_metadata_bucket(data) - suppressed: Final = _active_compression_guardrail_names() - ({policy.model} if policy.model else set()) + # Markers carry a per-process token so a caller cannot suppress a guardrail by + # naming it in its own request metadata. + suppressed: Final = tuple( + marker + for guardrail in _active_compression_guardrails() + if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) + ) if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = sorted(suppressed) + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = list(suppressed) if policy.model is not None: requested = metadata.get("guardrails") @@ -157,44 +181,52 @@ async def arm_pre_call(data: dict, llm_router: "Router | None") -> dict: snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) if snapshot is not None: - metadata[AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] = copy.deepcopy(snapshot) + _routing_messages_snapshot.set(tuple(dict(message) for message in snapshot)) return data +def _snapshot_messages() -> list[dict[str, Any]] | None: + snapshot: Final = _routing_messages_snapshot.get() + return None if snapshot is None else [dict(message) for message in snapshot] + + async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, messages: list[dict[str, Any]] | None, request_kwargs: Mapping[str, object], ) -> list[dict[str, Any]] | None: - """Messages to use for a routing decision, compressed per `policy.routing`. + """Messages to use for a routing decision, per `policy.routing`. - Returns None when there is no policy or the policy's routing side names no - compression, meaning the caller should route on whatever messages it already - has. The model call is untouched by this function either way: model-side - compression, if any, already ran as an ordinary pre-call guardrail before the - router was ever reached. + Returns None when the caller should route on whatever messages it already has. + The model call is untouched either way: model-side compression, if any, already + ran as an ordinary pre-call guardrail before the router was reached, so when the + two hops differ the routing decision reads the pre-compression snapshot rather + than what that guardrail left behind. """ - if policy is None or policy.routing is None: + if policy is None: + return None + + original: Final = _snapshot_messages() or messages + + if policy.routing is None: + # Explicitly no compression for routing. When the model side compressed, the + # messages in hand are its output, so fall back to the untouched snapshot. + return _snapshot_messages() if policy.model is not None else None + + if not original: return None from litellm.proxy.common_utils.registry_read_through import ( get_initialized_guardrail_with_read_through, ) - metadata_key: Final = get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.get(metadata_key) - snapshot: Final = metadata.get(AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY) if isinstance(metadata, dict) else None - original: Final = snapshot if isinstance(snapshot, list) else messages - if not original: - return None - guardrail: Final = await get_initialized_guardrail_with_read_through(policy.routing) if guardrail is None: verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return None + return original inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} # A throwaway request_data: apply_guardrail writes its stats onto this dict, not diff --git a/litellm/router.py b/litellm/router.py index 8149ec60ddc..bcb2e2aa7ff 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13039,11 +13039,19 @@ class Router: from litellm.proxy.guardrails.auto_router_compression import ( messages_for_routing, - policy_from_litellm_params, + policy_for_model, + team_id_from_request, ) - marker_params: Final = self._alias_marker_litellm_params(registered_model_name, selected_strategy.tags) - compression_policy: Final = policy_from_litellm_params(marker_params) if marker_params else None + # Resolved through the same tag-aware lookup the proxy's pre-call arming used, + # so an alias carrying several tag-scoped markers cannot suppress one marker's + # guardrail and then route under a different marker's policy. + compression_policy: Final = policy_for_model( + llm_router=self, + model_alias=registered_model_name, + team_id=team_id_from_request(request_kwargs), + request_tags=_get_tags_from_request_kwargs(request_kwargs), + ) # When both hops share the same compression, the model-side guardrail already # ran in the proxy's ordinary pre-call hook and compressed `messages` in place # (arm_pre_call armed it whether or not it is `default_on`); reuse that result @@ -13133,16 +13141,9 @@ class Router: return pre_routing_hook_response - def _alias_marker_litellm_params( + def _forwardable_alias_marker_params( self, model: str, strategy_tags: tuple[str, ...] - ) -> Mapping[str, object] | None: - """The auto-router marker deployment's own `litellm_params` for `model`, tag-scoped. - - Shared by `_forwardable_alias_marker_params` (forwarding api_base/api_key/... - gaps onto the routed deployment) and the auto-router compression policy lookup - (reading `auto_router_routing_compression`/`auto_router_model_compression`), so - both read the same marker row when an alias has more than one, tag-scoped marker. - """ + ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params for idx in self.model_name_to_deployment_indices.get(model, ()) @@ -13152,12 +13153,7 @@ class Router: tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags ) - return tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) - - def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] - ) -> tuple[tuple[str, object], ...]: - selected: Final = self._alias_marker_litellm_params(model, strategy_tags) + selected: Final = tag_matched[0] if tag_matched else (marker_params[0] if marker_params else None) if selected is None: return () return tuple( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 1fb4299cb56..f590903cb74 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -532,7 +532,9 @@ class TestCustomGuardrailShouldRunGuardrail: data = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["headroom-default"], + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + always_on.auto_router_suppression_marker() + ], }, } @@ -550,10 +552,13 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) + other = CustomGuardrail(guardrail_name="some-other-guardrail", default_on=True) data = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["some-other-guardrail"], + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + other.auto_router_suppression_marker() + ], }, } @@ -562,6 +567,32 @@ class TestCustomGuardrailShouldRunGuardrail: is True ) + def test_should_run_guardrail_ignores_a_forged_suppression_marker(self): + """A caller controls request metadata, so a bare guardrail name there must not + switch off an always-on guardrail: only the per-process marker counts.""" + from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.types.guardrails import GuardrailEventHooks + + always_on = CustomGuardrail( + guardrail_name="headroom-default", + default_on=True, + event_hook=GuardrailEventHooks.pre_call, + ) + forged = { + "model": "smart-router", + "metadata": { + AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + "headroom-default", + "forged-token:headroom-default", + ], + }, + } + + assert ( + always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) + is True + ) + class TestApplyGuardrailCheck: def test_apply_guardrail_check_only_on_direct_implementation(self): diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index b2e83e75768..b906e60bb86 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -4,28 +4,33 @@ Unit tests for litellm.proxy.guardrails.auto_router_compression. Covers: - policy_from_litellm_params: absent keys mean no policy; the "none" sentinel normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias -- arm_pre_call: no-op without a policy; suppresses active compression guardrails; - arms the model-side guardrail even when it isn't default_on; snapshots messages -- messages_for_routing: no-op without a policy or an unset routing side; compresses - via the named guardrail's apply_guardrail; never writes stats onto the caller's - own request_kwargs (regression for double-counted compression savings) +- policy_for_model: finds the auto-router marker deployment for an alias, and + picks the tag-scoped marker the request's tags actually match +- arm_pre_call: no-op without a policy; suppresses active compression guardrails + with a forgery-proof marker; arms the model-side guardrail even when it isn't + default_on; keeps the pre-compression snapshot out of persisted metadata +- messages_for_routing: no-op without a policy; routes on the pre-compression + snapshot when the two hops differ; compresses via the named guardrail's + apply_guardrail; never writes stats onto the caller's own request_kwargs + (regression for double-counted compression savings) """ +import json from typing import Any import pytest +from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy.guardrails import auto_router_compression from litellm.proxy.guardrails.auto_router_compression import ( - AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY, AutoRouterCompressionPolicy, arm_pre_call, messages_for_routing, policy_for_model, policy_from_litellm_params, ) -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs @@ -76,36 +81,61 @@ class _FakeRouter: return [d for d in self._deployments if d.get("model_name") == model_name] +def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: + return { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + **compression, + **({"tags": tags} if tags is not None else {}), + }, + } + + class TestPolicyForModel: def test_no_router_returns_none(self): - assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None) is None + assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None def test_no_marker_deployment_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None def test_marker_deployment_with_policy_is_found(self): + router = _FakeRouter( + [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + + def test_picks_the_marker_whose_tags_the_request_carries(self): + """Regression: an alias with several tag-scoped markers must not suppress one + marker's guardrail and then route under a different marker's policy.""" router = _FakeRouter( [ - { - "model_name": "smart-router", - "litellm_params": { - "model": "auto_router/complexity_router", - "auto_router_routing_compression": "headroom-a", - "auto_router_model_compression": "none", - }, - } + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-us"}, tags=["us"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None) + + eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + + assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) + assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) + + def test_untagged_marker_matches_any_request(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + policy = policy_for_model( + llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) @@ -186,10 +216,25 @@ class TestArmPreCall: data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} result = await arm_pre_call(data=data, llm_router=router) suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert "always-on-compression" in suppressed + assert suppressed == [always_on.auto_router_suppression_marker()] + # The bare name alone must never suppress: that is what a caller could forge. + assert "always-on-compression" not in suppressed + assert always_on.should_run_guardrail(data=result, event_type=GuardrailEventHooks.pre_call) is False finally: litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) + @pytest.mark.asyncio + async def test_a_caller_cannot_suppress_a_guardrail_by_naming_it_in_metadata(self): + """Regression: request metadata is caller-controlled, so a bare guardrail name + there must not switch off a PII, content-filter, or compression guardrail.""" + guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") + forged = { + "model": "smart-router", + "metadata": {AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["always-on-compression"]}, + } + + assert guardrail._suppressed_by_auto_router_compression(forged) is False + @pytest.mark.asyncio async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): router = _FakeRouter( @@ -209,43 +254,82 @@ class TestArmPreCall: assert result["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio - async def test_snapshots_original_messages(self): - router = _FakeRouter( - [ - { - "model_name": "smart-router", - "litellm_params": { - "model": "auto_router/complexity_router", - "auto_router_routing_compression": "headroom-a", - "auto_router_model_compression": "none", - }, - } - ] - ) - original_messages = [{"role": "user", "content": "hi"}] + async def test_snapshot_never_lands_in_persisted_metadata(self): + """Regression: refresh_proxy_server_request_body_snapshot copies metadata into + proxy_server_request.body, which deployments persist to spend logs. The + pre-compression snapshot holds the prompt before any masking guardrail ran, so + it must live outside anything that gets serialized.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] data = {"model": "smart-router", "messages": original_messages} + result = await arm_pre_call(data=data, llm_router=router) - snapshot = result["metadata"][AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY] - assert snapshot == original_messages - assert snapshot is not original_messages # a copy, not the live reference + + assert "123-45-6789" not in json.dumps(result["metadata"]) + assert auto_router_compression._snapshot_messages() == original_messages + + @pytest.mark.asyncio + async def test_snapshot_is_a_copy_not_the_live_message_list(self): + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + original_messages = [{"role": "user", "content": "hi"}] + + await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) + original_messages[0]["content"] = "mutated after the snapshot" + + assert auto_router_compression._snapshot_messages() == [{"role": "user", "content": "hi"}] + + @pytest.mark.asyncio + async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): + router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + await arm_pre_call(data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, + llm_router=router_with) + + router_without = _FakeRouter( + [{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + ) + await arm_pre_call(data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, + llm_router=router_without) + + assert auto_router_compression._snapshot_messages() is None class TestMessagesForRouting: + @pytest.fixture(autouse=True) + def _clear_snapshot(self): + auto_router_compression._routing_messages_snapshot.set(None) + yield + auto_router_compression._routing_messages_snapshot.set(None) + @pytest.mark.asyncio async def test_no_policy_returns_none(self): assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_routing_side_unset_returns_none(self): - policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + async def test_routing_none_with_no_model_compression_returns_none(self): + """Nothing compressed either hop, so the caller's own messages are already right.""" + policy = AutoRouterCompressionPolicy(routing=None, model=None) assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_unknown_guardrail_name_returns_none(self): + async def test_routing_none_with_model_compression_routes_on_the_snapshot(self): + """Regression: with routing explicitly off and the model side compressed, the + messages in hand are the model-side guardrail's output. Routing asked for no + compression, so it must read the pre-compression snapshot instead.""" + original = [{"role": "user", "content": "the full original conversation"}] + auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original)) + policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") + model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] + + result = await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) + + assert result == original + + @pytest.mark.asyncio + async def test_unknown_guardrail_name_routes_on_the_uncompressed_messages(self): policy = AutoRouterCompressionPolicy(routing="does-not-exist", model=None) messages = [{"role": "user", "content": "hi"}] result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) - assert result is None + assert result == messages @pytest.mark.asyncio async def test_compresses_via_the_named_guardrail(self, registered_guardrail): @@ -256,15 +340,14 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_uses_the_snapshot_when_present(self, registered_guardrail): - policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) - snapshot = [{"role": "user", "content": "original"}] - request_kwargs = {"metadata": {AUTO_ROUTER_ROUTING_MESSAGES_SNAPSHOT_KEY: snapshot}} - # `messages` here stands in for whatever a model-side guardrail already + policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") + auto_router_compression._routing_messages_snapshot.set(({"role": "user", "content": "original"},)) + # `messages` here stands in for whatever the model-side guardrail already # rewrote `data["messages"]` to -- routing must ignore it and compress the # pristine snapshot instead. already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] result = await messages_for_routing( - policy=policy, messages=already_rewritten, request_kwargs=request_kwargs + policy=policy, messages=already_rewritten, request_kwargs={} ) assert result == [{"role": "user", "content": "[COMPRESSED] original"}] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cdae3b131ae..4c5813911ac 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10105,6 +10105,35 @@ class TestAutoRouterCompressionDecoupling: assert response.messages == original_messages assert registered_guardrail.call_count == 0 + @pytest.mark.asyncio + async def test_routing_none_routes_on_the_original_not_the_model_compressed_messages( + self, registered_guardrail + ): + """Regression: with routing explicitly off and the model side compressed, the + messages the router holds are the model-side guardrail's output. Routing asked + for no compression, so it has to classify on the pre-compression snapshot.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "none", + "auto_router_model_compression": "fake-compress", + } + ) + original_messages = self._messages() + auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original_messages)) + model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + + try: + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) + finally: + auto_router_compression._routing_messages_snapshot.set(None) + + assert strategy.received_messages == original_messages + assert registered_guardrail.call_count == 0 + @pytest.mark.asyncio async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): """The same/different distinction exists so a shared choice does not pay for diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index a0a240f76b0..52ea7645034 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -5,8 +5,7 @@ import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Info } from "lucide-react"; import React from "react"; import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; -import { COMPRESSION_GUARDRAIL_PROVIDER } from "@/app/(dashboard)/cost-optimization/_components/helpers"; -import { NO_COMPRESSION } from "./buildAutoRouterCompression"; +import { isCompressionGuardrailProvider, NO_COMPRESSION } from "./buildAutoRouterCompression"; interface CompressionControlsProps { routing: string | undefined; @@ -29,7 +28,7 @@ const CompressionControls: React.FC = ({ }) => { const { data } = useGuardrails(); const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) - .filter((g) => (g.litellm_params?.guardrail ?? "").toString().toLowerCase() === COMPRESSION_GUARDRAIL_PROVIDER) + .filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail)) .map((g) => ({ label: g.guardrail_name, value: g.guardrail_name })); const options: SearchSelectOption[] = [NONE_OPTION, ...compressionOptions]; diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 49180d5b4ce..5afdcf2b15c 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -12,6 +12,13 @@ export const NO_COMPRESSION = "none"; +/** Guardrail providers that compress prompts, mirroring COMPRESSION_GUARDRAIL_PROVIDERS in + * litellm/proxy/guardrails/auto_router_compression.py. Both are selectable per hop. */ +export const COMPRESSION_GUARDRAIL_PROVIDERS: readonly string[] = ["headroom", "compresr"]; + +export const isCompressionGuardrailProvider = (provider: unknown): boolean => + typeof provider === "string" && COMPRESSION_GUARDRAIL_PROVIDERS.includes(provider.toLowerCase()); + export interface AutoRouterCompressionState { routing: string | undefined; sameAsRouting: boolean; From 5980055d7eae3d1ca28286979c5bd264cd37af57 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 16:50:35 -0700 Subject: [PATCH 12/61] feat(shadow_eval): say which shape produced an unparseable judge verdict The parser message alone cannot separate a judge that answered with nothing from one truncated mid-object, and the two want opposite fixes. Records the reply's shape, never its text, since no attempt row carries sampled content. --- litellm/integrations/shadow_eval_logger.py | 20 +++- .../integrations/test_shadow_eval_logger.py | 94 +++++++++++++++++++ 2 files changed, 113 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index fc82ebafe09..fb75ef74db9 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -345,6 +345,22 @@ def _failure_detail(e: BaseException) -> str: return f"{type(e).__name__}{location}: {e}" +def _judge_reply_shape(response: object) -> str: + """How an unparseable judge reply was shaped. The parser's own message cannot separate a + judge that answered with nothing from one truncated mid-object, and those want opposite + fixes. Shape only, never the reply text: the judge quotes the sampled turns it compares, + and no attempt row carries sampled content today.""" + try: + choice: Final = response["choices"][0] # pyright: ignore[reportIndexIssue] # judge replies are subscriptable payloads + content: Final = choice["message"]["content"] + finish: Final = choice.get("finish_reason") or "unknown" + except (AttributeError, KeyError, IndexError, TypeError): + return "unreadable judge reply" + served: Final = str(getattr(response, "model", None) or "unknown") + body: Final = f"{len(str(content))} chars" if content else "no content" + return f"finish_reason={finish}, content={body}, model={served}" + + def _call_cost(response: object) -> float: """Price one eval-arm call with the figure the spend pipeline bills: the router client stamps _hidden_params.response_cost from the deployment's own pricing, which the public @@ -1139,7 +1155,9 @@ class ShadowEvalLogger(CustomLogger): verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw)) except Exception as e: # noqa: BLE001 # malformed verdicts become error rows verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e) - return _CallFailure(f"unparseable judge verdict: {e}", cost=_call_cost(response)) + return _CallFailure( + f"unparseable judge verdict: {e}; {_judge_reply_shape(response)}", cost=_call_cost(response) + ) return _JudgeVerdict( preference=_unmask_preference(verdict.preference, real_is_a), confidence=max(0.0, min(1.0, verdict.confidence)), diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 9fcbd116f63..dbc6d4ec915 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -141,6 +141,27 @@ def _reasoning_judge_router( return router +def _judge_reply_router(content: str | None, finish_reason: str = "stop", served_model: str = "judge-pick") -> MagicMock: + """A router whose judge arm returns a caller-shaped reply, so the shapes that all land + on the same parser error can be posed apart: no content at all, versus JSON cut off + mid-object.""" + router = MagicMock() + router.model_group_alias = {} + router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + async def acompletion(**kwargs): + if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: + kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} + return {"choices": [{"message": {"content": "shadow answer"}}]} + return ModelResponse( + model=served_model, + choices=[{"index": 0, "finish_reason": finish_reason, "message": {"role": "assistant", "content": content}}], + ) + + router.acompletion = MagicMock(side_effect=acompletion) + return router + + def _spend_counter(store=None): """In-memory stand-in for the proxy's cross-pod spend counter: reads take the max of the counter and the caller's fallback, exactly like get_current_spend does for a key @@ -1126,6 +1147,79 @@ class TestShadowPipeline: assert row["judge_cost"] == expected_cost assert row["shadow_cost"] == expected_shadow_cost + async def _judge_error(self, router: MagicMock, monkeypatch: pytest.MonkeyPatch) -> str: + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "completion_cost", lambda completion_response: 0.007) + prisma = _prisma() + await _logger(router=router, prisma=prisma)._run_shadow_eval( + job=_job(), + request_id="req-1", + messages=({"role": "user", "content": "hi"},), + real_text="real answer", + real_model="claude-opus", + real_cost=0.0, + real_classifier_cost=0.0, + real_cache_hit=False, + control_tier=None, + shadow_params={}, + parent_metadata={}, + ) + return prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"]["error"] + + async def test_a_judge_that_answered_nothing_is_told_apart_from_one_cut_off( + self, monkeypatch: pytest.MonkeyPatch + ): + """Both land on the same parser message, and they want opposite fixes: a judge + returning no content points at the reply never being text, while one cut off + mid-object points at the output cap. The row has to say which.""" + truncated = '{"preference": "A", "confidence": 0.9, "reasoning": "' + answered_nothing = await self._judge_error(_judge_reply_router(None), monkeypatch) + cut_off = await self._judge_error( + _judge_reply_router(truncated, finish_reason="length"), monkeypatch + ) + + assert "content=no content" in answered_nothing + assert "finish_reason=stop" in answered_nothing + assert f"content={len(truncated)} chars" in cut_off + assert "finish_reason=length" in cut_off + + async def test_an_unparseable_verdict_names_the_model_that_served_it(self, monkeypatch: pytest.MonkeyPatch): + """A judge_model that fans out over deployments hides which one truncates: without + the served model the operator cannot tell a bad deployment from a bad cap.""" + error = await self._judge_error(_judge_reply_router(None, served_model="claude-sonnet-5"), monkeypatch) + + assert "model=claude-sonnet-5" in error + + async def test_a_diagnosed_verdict_error_stays_groupable(self, monkeypatch: pytest.MonkeyPatch): + """The customer groups attempt rows by error text. Every varying part has to sit + after the first semicolon or each row becomes its own group.""" + first = await self._judge_error(_judge_reply_router(None, served_model="model-a"), monkeypatch) + second = await self._judge_error(_judge_reply_router(None, served_model="model-b"), monkeypatch) + + assert first != second + assert first.split(";")[0] == second.split(";")[0] + + async def test_a_judge_reply_that_cannot_be_read_still_records_an_error(self, monkeypatch: pytest.MonkeyPatch): + """The shape reader runs inside the failure path: it must never raise a second time + and cost the row entirely.""" + router = MagicMock() + router.model_group_alias = {} + router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}]) + + async def acompletion(**kwargs): + if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: + kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} + return {"choices": [{"message": {"content": "shadow answer"}}]} + return {"choices": []} + + router.acompletion = MagicMock(side_effect=acompletion) + + error = await self._judge_error(router, monkeypatch) + + assert "unparseable judge verdict" in error + assert "unreadable judge reply" in error + async def test_an_empty_shadow_reply_still_bills_its_cost(self, monkeypatch: pytest.MonkeyPatch): """A shadow call that returns no extractable text has still billed; pricing it at zero would keep the dollar gate open while shadow calls keep charging the key.""" From d0a80067377cb6978deac35492b6681a65c991fc Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 17:03:01 -0700 Subject: [PATCH 13/61] fix(auto-router compression): tag-scoped markers now take precedence over untagged An untagged marker (no tags key or empty tags list) was matching every request because requested.issuperset(frozenset()) is always true. When an alias carried multiple markers, the loop tried tag-matched markers first, but an untagged one could still match the tag-match query, and then the first one with a policy would be returned. Now only markers with a non-empty tags list can match via the tag-specific lookup; untagged markers are tried only after all tag-specific ones. Regression test added: test_tag_scoped_marker_takes_precedence_over_untagged fails with the old code. Also removed unused Any import per greptile's typing note. --- .../proxy/guardrails/auto_router_compression.py | 16 ++++++++++------ .../guardrails/test_auto_router_compression.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 7ccd1937543..490ce550003 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -15,7 +15,7 @@ each hop sees. import contextvars from collections.abc import Mapping, MutableMapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY @@ -93,7 +93,9 @@ def policy_for_model( ) requested: Final = frozenset(request_tags) tag_matched: Final = tuple( - params for params in markers if requested.issuperset(frozenset(params.get("tags") or ())) + params + for params in markers + if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) for params in (*tag_matched, *markers): policy = policy_from_litellm_params(params) @@ -186,16 +188,16 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | return data -def _snapshot_messages() -> list[dict[str, Any]] | None: +def _snapshot_messages() -> list[dict[str, object]] | None: snapshot: Final = _routing_messages_snapshot.get() return None if snapshot is None else [dict(message) for message in snapshot] async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, - messages: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, request_kwargs: Mapping[str, object], -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """Messages to use for a routing decision, per `policy.routing`. Returns None when the caller should route on whatever messages it already has. @@ -228,7 +230,9 @@ async def messages_for_routing( ) return original - inputs: GenericGuardrailAPIInputs = {"structured_messages": [dict(m) for m in original]} + inputs: GenericGuardrailAPIInputs = { + "structured_messages": [dict(m) for m in original] # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape + } # A throwaway request_data: apply_guardrail writes its stats onto this dict, not # the real request's metadata, so routing-side compression never double-counts # against extract_compression_saved_tokens's model-savings accounting. diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index b906e60bb86..79f069d8a2e 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -138,6 +138,18 @@ class TestPolicyForModel: ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + def test_tag_scoped_marker_takes_precedence_over_untagged(self): + """Regression: when multiple markers exist, the tag-scoped one the request + actually matches should be used, not the first untagged one.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-untagged"}), + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) + class _RecordingCompressionGuardrail(CustomGuardrail): """A guardrail whose apply_guardrail marks every text message as compressed.""" From e273cf301fc710a88bb3820e3d35f7fe6a6b20bb Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 17:21:04 -0700 Subject: [PATCH 14/61] fix(ci): satisfy ruff format, prettier, and eslint max-lines gates - ruff format on auto_router_compression.py (a long comprehension wrapped across three lines instead of one) - prettier on buildAutoRouterCompression.ts and the two test files it touched - ComplexityRouterConfig.tsx crossed the 800-line eslint max-lines ceiling once the compression accordion entry landed. Extracted TierRowSelect into its own file (already self-contained, used only within this file and PlanModeOverrideControls) and simplified CompressionControls' props to a single state/onChange pair instead of six individual callbacks, moving the per-field derivation into the component that already owns this state shape --- .../guardrails/auto_router_compression.py | 4 +- .../add_model/ComplexityRouterConfig.tsx | 40 +------------------ .../add_model/CompressionControls.tsx | 29 +++++++------- .../components/add_model/TierRowSelect.tsx | 25 ++++++++++++ .../add_model/buildAutoRouterCompression.ts | 2 +- .../edit_auto_router_modal.test.tsx | 7 ++-- 6 files changed, 47 insertions(+), 60 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 490ce550003..2122d353def 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -93,9 +93,7 @@ def policy_for_model( ) requested: Final = frozenset(request_tags) tag_matched: Final = tuple( - params - for params in markers - if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) + params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) for params in (*tag_matched, *markers): policy = policy_from_litellm_params(params) diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index 42115265034..92fd9893995 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,11 +1,11 @@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { SearchSelect } from "@/components/shared/SearchSelect"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { ChevronRight, Info, Plus, Trash2, X } from "lucide-react"; import { Switch } from "@/components/ui/switch"; import { AffinityControls } from "./AffinityControls"; +import TierRowSelect from "./TierRowSelect"; import { ModalityRoutingControls } from "./ModalityRoutingControls"; import { Card, CardContent } from "@/components/ui/card"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; @@ -370,27 +370,6 @@ const TierRowEditFields: React.FC<{ ); -const TierRowSelect: React.FC<{ - label: string; - options: { value: string; label: string }[]; - value: string | null; - onValueChange: (rowId: string) => void; - placeholder?: string; -}> = ({ label, options, value, onValueChange, placeholder }) => ( - -); - export type AdaptiveEligible = "all" | "classified_tier"; export type ComplexityTierLabels = Partial>; @@ -889,22 +868,7 @@ const ComplexityRouterConfig: React.FC = ({ key: "compression", label: Advanced: Compression, children: ( - - onAutoRouterCompressionChange({ - ...autoRouterCompression, - routing, - sameAsRouting: routing === undefined ? true : autoRouterCompression.sameAsRouting, - }) - } - sameAsRouting={autoRouterCompression.sameAsRouting} - onSameAsRoutingChange={(sameAsRouting) => - onAutoRouterCompressionChange({ ...autoRouterCompression, sameAsRouting }) - } - model={autoRouterCompression.model} - onModelChange={(model) => onAutoRouterCompressionChange({ ...autoRouterCompression, model })} - /> + ), }, ] diff --git a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx index 52ea7645034..c1817918f60 100644 --- a/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/CompressionControls.tsx @@ -5,27 +5,26 @@ import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; import { Info } from "lucide-react"; import React from "react"; import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; -import { isCompressionGuardrailProvider, NO_COMPRESSION } from "./buildAutoRouterCompression"; +import { + AutoRouterCompressionState, + isCompressionGuardrailProvider, + NO_COMPRESSION, +} from "./buildAutoRouterCompression"; interface CompressionControlsProps { - routing: string | undefined; - onRoutingChange: (value: string | undefined) => void; - sameAsRouting: boolean; - onSameAsRoutingChange: (same: boolean) => void; - model: string | undefined; - onModelChange: (value: string | undefined) => void; + value: AutoRouterCompressionState; + onChange: (state: AutoRouterCompressionState) => void; } const NONE_OPTION: SearchSelectOption = { label: "None (no compression)", value: NO_COMPRESSION }; -const CompressionControls: React.FC = ({ - routing, - onRoutingChange, - sameAsRouting, - onSameAsRoutingChange, - model, - onModelChange, -}) => { +const CompressionControls: React.FC = ({ value, onChange }) => { + const { routing, sameAsRouting, model } = value; + const onRoutingChange = (newRouting: string | undefined) => + onChange({ ...value, routing: newRouting, sameAsRouting: newRouting === undefined ? true : sameAsRouting }); + const onSameAsRoutingChange = (newSameAsRouting: boolean) => onChange({ ...value, sameAsRouting: newSameAsRouting }); + const onModelChange = (newModel: string | undefined) => onChange({ ...value, model: newModel }); + const { data } = useGuardrails(); const compressionOptions: SearchSelectOption[] = (data?.guardrails ?? []) .filter((g) => isCompressionGuardrailProvider(g.litellm_params?.guardrail)) diff --git a/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx b/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx new file mode 100644 index 00000000000..ad7d53f9eae --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/TierRowSelect.tsx @@ -0,0 +1,25 @@ +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import React from "react"; + +const TierRowSelect: React.FC<{ + label: string; + options: { value: string; label: string }[]; + value: string | null; + onValueChange: (rowId: string) => void; + placeholder?: string; +}> = ({ label, options, value, onValueChange, placeholder }) => ( + +); + +export default TierRowSelect; diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 5afdcf2b15c..c86416b507f 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -42,7 +42,7 @@ export const buildAutoRouterCompressionParams = ( if (state.routing === undefined) return {}; return { auto_router_routing_compression: state.routing, - auto_router_model_compression: state.sameAsRouting ? state.routing : (state.model ?? NO_COMPRESSION), + auto_router_model_compression: state.sameAsRouting ? state.routing : state.model ?? NO_COMPRESSION, }; }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index b93db8d963e..d8474b492e1 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -1040,9 +1040,10 @@ describe("EditAutoRouterModal prompt compression", () => { return payload?.litellm_params; }; - const renderWithStoredCompression = ( - compression?: { auto_router_routing_compression?: string; auto_router_model_compression?: string }, - ) => + const renderWithStoredCompression = (compression?: { + auto_router_routing_compression?: string; + auto_router_model_compression?: string; + }) => renderWithProviders( Date: Fri, 4 Sep 2026 17:42:52 -0700 Subject: [PATCH 15/61] refactor(auto-router compression): satisfy the LIT001/LIT002 type-discipline gate The gate has no headroom, so the new module had to stop introducing mutable collections rather than spend budget on them: - the marker lookup falls back to () and drops an `or {}` that isinstance already covered - the suppression list is stored as the tuple it was built as; the read side in custom_guardrail accepts list or tuple, since JSON round-trips it to a list - the snapshot holds MappingProxyType entries, so it is immutable at rest and _snapshot_messages can hand back the stored tuple with no defensive copy - arm_pre_call returns None instead of echoing back the dict it mutates in place - _suppressed_by_auto_router_compression takes a Mapping, which is all it reads The four remaining mutable spots are external contracts, each suppressed with the reason: the pre-routing hook protocol types messages as list[dict], the metadata["guardrails"] key is extended by litellm_pre_call_utils via an isinstance(..., list) check, apply_guardrail takes a dict it writes stats into, and pydantic's model_copy takes a dict. --- litellm/integrations/custom_guardrail.py | 8 +- litellm/proxy/common_request_processing.py | 2 +- .../guardrails/auto_router_compression.py | 85 +++++++++++-------- litellm/router.py | 3 +- .../test_auto_router_compression.py | 65 ++++++-------- 5 files changed, 83 insertions(+), 80 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 558e97cfc16..1ec08641706 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -954,16 +954,18 @@ class CustomGuardrail(CustomLogger): return None return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - def _suppressed_by_auto_router_compression(self, data: dict[str, object]) -> bool: + def _suppressed_by_auto_router_compression(self, data: Mapping[str, object]) -> bool: """True when an auto router's own compression policy suppresses this guardrail.""" marker: Final = self.auto_router_suppression_marker() if marker is None: return False for meta_key in ("metadata", "litellm_metadata"): meta = data.get(meta_key) - if isinstance(meta, dict): + if isinstance(meta, Mapping): + # arm_pre_call writes a tuple; it arrives as a list once the metadata + # has been round-tripped through JSON. suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, list) and marker in suppressed: + if isinstance(suppressed, (list, tuple)) and marker in suppressed: return True return False diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 534b2db3e61..5e6c9b34332 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2009,7 +2009,7 @@ class ProxyBaseLLMRequestProcessing: # request: suppress every other compression guardrail and arm whichever one # the policy names for the model call, before those guardrails get a chance # to run below. - self.data = await _arm_auto_router_compression(data=self.data, llm_router=llm_router) + await _arm_auto_router_compression(data=self.data, llm_router=llm_router) self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 2122d353def..d26479f7de0 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -13,8 +13,9 @@ each hop sees. """ import contextvars -from collections.abc import Mapping, MutableMapping, Sequence +from collections.abc import Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger @@ -84,11 +85,11 @@ def policy_for_model( """ if llm_router is None: return None - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or [] + deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () markers: Final = tuple( litellm_params for deployment in deployments - if isinstance(litellm_params := deployment.get("litellm_params") or {}, Mapping) + if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) ) requested: Final = frozenset(request_tags) @@ -128,7 +129,10 @@ def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name) -async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | None") -> MutableMapping[str, object]: +async def arm_pre_call( + data: MutableMapping[str, object], # mutable-ok: arms the live request dict in place + llm_router: "Router | None", +) -> None: """Apply an auto router's compression policy, if any, before guardrails run. Suppresses every other compression guardrail, re-enables the model-side @@ -138,11 +142,11 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | """ _routing_messages_snapshot.set(None) if llm_router is None: - return data + return model_alias: Final = data.get("model") if not isinstance(model_alias, str) or not model_alias: - return data + return # Read-only until a policy is confirmed: creating the metadata bucket for every # request, including the vast majority with no auto-router compression policy, @@ -156,7 +160,7 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | request_tags=_get_tags_from_request_kwargs(data), ) if policy is None: - return data + return _, metadata = get_or_create_metadata_bucket(data) # Markers carry a per-process token so a caller cannot suppress a guardrail by @@ -167,35 +171,41 @@ async def arm_pre_call(data: MutableMapping[str, object], llm_router: "Router | if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) ) if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = list(suppressed) + metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = suppressed if policy.model is not None: - requested = metadata.get("guardrails") - if isinstance(requested, list): - if policy.model not in requested: - requested.append(policy.model) - else: - metadata["guardrails"] = [policy.model] + requested: Final = metadata.get("guardrails") + existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () + if policy.model not in existing: + # A list, not a tuple: litellm_pre_call_utils tests this key with + # isinstance(..., list) and extends it, and would drop a tuple on the floor. + metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) if snapshot is not None: - _routing_messages_snapshot.set(tuple(dict(message) for message in snapshot)) - - return data + _routing_messages_snapshot.set(tuple(MappingProxyType(dict(message)) for message in snapshot)) -def _snapshot_messages() -> list[dict[str, object]] | None: - snapshot: Final = _routing_messages_snapshot.get() - return None if snapshot is None else [dict(message) for message in snapshot] +def _snapshot_messages() -> tuple[Mapping[str, object], ...] | None: + return _routing_messages_snapshot.get() + + +def _as_routing_messages( + messages: Iterable[Mapping[str, object]], +) -> list[dict[str, object]]: # mutable-ok: shape fixed by the pre-routing hook protocol + """A fresh, independently mutable copy, the shape the pre-routing hook takes.""" + return [dict(message) for message in messages] # mutable-ok: shape fixed by the pre-routing hook protocol async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, - messages: list[dict[str, object]] | None, + # list[dict], not Sequence[Mapping]: the async_pre_routing_hook protocol in + # litellm/types/router.py types `messages` as list[dict[str, Any]]. + messages: list[dict[str, object]] | None, # mutable-ok: shape fixed by the pre-routing hook protocol request_kwargs: Mapping[str, object], -) -> list[dict[str, object]] | None: +) -> list[dict[str, object]] | None: # mutable-ok: shape fixed by the pre-routing hook protocol """Messages to use for a routing decision, per `policy.routing`. Returns None when the caller should route on whatever messages it already has. @@ -207,12 +217,13 @@ async def messages_for_routing( if policy is None: return None - original: Final = _snapshot_messages() or messages + snapshot: Final = _snapshot_messages() + original: Final = snapshot if snapshot is not None else messages if policy.routing is None: # Explicitly no compression for routing. When the model side compressed, the # messages in hand are its output, so fall back to the untouched snapshot. - return _snapshot_messages() if policy.model is not None else None + return _as_routing_messages(snapshot) if policy.model is not None and snapshot is not None else None if not original: return None @@ -226,20 +237,20 @@ async def messages_for_routing( verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return original + return _as_routing_messages(original) - inputs: GenericGuardrailAPIInputs = { - "structured_messages": [dict(m) for m in original] # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape - } - # A throwaway request_data: apply_guardrail writes its stats onto this dict, not - # the real request's metadata, so routing-side compression never double-counts - # against extract_compression_saved_tokens's model-savings accounting. - throwaway_request_data: Final[dict[str, object]] = { - "messages": original, - "model": request_kwargs.get("model"), + inputs: Final[GenericGuardrailAPIInputs] = { + "structured_messages": _as_routing_messages(original) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape } + model: Final = request_kwargs.get("model") + # A throwaway request_data: apply_guardrail writes its stats onto this dict, not the + # real request's metadata, so routing-side compression never double-counts against + # extract_compression_saved_tokens's model-savings accounting. + stats_sink: Final = {"messages": original, "model": model} # mutable-ok: apply_guardrail writes its stats here result: Final = await guardrail.apply_guardrail( - inputs=inputs, request_data=throwaway_request_data, input_type="request" + inputs=inputs, + request_data=stats_sink, + input_type="request", ) - compressed = result.get("structured_messages") - return compressed if isinstance(compressed, list) else original + compressed: Final = result.get("structured_messages") + return compressed if isinstance(compressed, list) else _as_routing_messages(original) diff --git a/litellm/router.py b/litellm/router.py index bcb2e2aa7ff..9637ad98c9a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13084,7 +13084,8 @@ class Router: and routing_messages is not None and pre_routing_hook_response.messages == routing_messages ): - pre_routing_hook_response = pre_routing_hook_response.model_copy(update={"messages": messages}) + restored: Final = {"messages": messages} # mutable-ok: pydantic's model_copy takes a dict + pre_routing_hook_response = pre_routing_hook_response.model_copy(update=restored) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 79f069d8a2e..de667e8ed48 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -97,9 +97,7 @@ class TestPolicyForModel: assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None def test_no_marker_deployment_returns_none(self): - router = _FakeRouter( - [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] - ) + router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): @@ -163,9 +161,7 @@ class _RecordingCompressionGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: self.request_data_seen.append(request_data) structured_messages = inputs.get("structured_messages") or [] - compressed = [ - {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages - ] + compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages] return {**inputs, "structured_messages": compressed} @@ -183,19 +179,16 @@ class TestArmPreCall: @pytest.mark.asyncio async def test_no_router_is_noop(self): data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=None) - assert result == data - assert "metadata" not in result + await arm_pre_call(data=data, llm_router=None) + assert "metadata" not in data @pytest.mark.asyncio async def test_no_policy_does_not_create_metadata_bucket(self): - router = _FakeRouter( - [{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}] - ) + router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=router) - assert "metadata" not in result - assert "litellm_metadata" not in result + await arm_pre_call(data=data, llm_router=router) + assert "metadata" not in data + assert "litellm_metadata" not in data @pytest.mark.asyncio async def test_policy_suppresses_active_compression_guardrails(self, monkeypatch): @@ -226,12 +219,12 @@ class TestArmPreCall: ] ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=router) - suppressed = result["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert suppressed == [always_on.auto_router_suppression_marker()] + await arm_pre_call(data=data, llm_router=router) + suppressed = data["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] + assert tuple(suppressed) == (always_on.auto_router_suppression_marker(),) # The bare name alone must never suppress: that is what a caller could forge. assert "always-on-compression" not in suppressed - assert always_on.should_run_guardrail(data=result, event_type=GuardrailEventHooks.pre_call) is False + assert always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is False finally: litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) @@ -262,8 +255,8 @@ class TestArmPreCall: ] ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - result = await arm_pre_call(data=data, llm_router=router) - assert result["metadata"]["guardrails"] == ["headroom-b"] + await arm_pre_call(data=data, llm_router=router) + assert data["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio async def test_snapshot_never_lands_in_persisted_metadata(self): @@ -275,10 +268,10 @@ class TestArmPreCall: original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] data = {"model": "smart-router", "messages": original_messages} - result = await arm_pre_call(data=data, llm_router=router) + await arm_pre_call(data=data, llm_router=router) - assert "123-45-6789" not in json.dumps(result["metadata"]) - assert auto_router_compression._snapshot_messages() == original_messages + assert "123-45-6789" not in json.dumps(data["metadata"]) + assert [dict(m) for m in auto_router_compression._snapshot_messages()] == original_messages @pytest.mark.asyncio async def test_snapshot_is_a_copy_not_the_live_message_list(self): @@ -288,19 +281,19 @@ class TestArmPreCall: await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) original_messages[0]["content"] = "mutated after the snapshot" - assert auto_router_compression._snapshot_messages() == [{"role": "user", "content": "hi"}] + assert [dict(m) for m in auto_router_compression._snapshot_messages()] == [{"role": "user", "content": "hi"}] @pytest.mark.asyncio async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - await arm_pre_call(data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, - llm_router=router_with) - - router_without = _FakeRouter( - [{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}] + await arm_pre_call( + data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, llm_router=router_with + ) + + router_without = _FakeRouter([{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + await arm_pre_call( + data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, llm_router=router_without ) - await arm_pre_call(data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, - llm_router=router_without) assert auto_router_compression._snapshot_messages() is None @@ -358,15 +351,11 @@ class TestMessagesForRouting: # rewrote `data["messages"]` to -- routing must ignore it and compress the # pristine snapshot instead. already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] - result = await messages_for_routing( - policy=policy, messages=already_rewritten, request_kwargs={} - ) + result = await messages_for_routing(policy=policy, messages=already_rewritten, request_kwargs={}) assert result == [{"role": "user", "content": "[COMPRESSED] original"}] @pytest.mark.asyncio - async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs( - self, registered_guardrail - ): + async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): """Regression: a real compression guardrail writes its stats onto whatever `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ request_data`). If that were the caller's own `request_kwargs`, routing-side From 88ada40cdad9e711e7b40d5d174464667474585c Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 18:37:38 -0700 Subject: [PATCH 16/61] fix(type-checking): satisfy the basedpyright budget gate for auto-router compression Two fixes for the zero-headroom basedpyright budget: - arm_pre_call's data parameter is dict[str, object], not MutableMapping: the latter is itself banned by LIT001 with no benefit, and it mismatched every dict-typed helper (get_or_create_metadata_bucket, resolve_structured_messages, _get_tags_from_request_kwargs), which is what the budget was actually flagging. - Router.async_pre_routing_hook computed pre_routing_hook_response in one shot instead of reassigning a Final-annotated local. The remaining two reportArgumentType hits are pre-existing: LiteLLM_Params(**merged) in _create_deployment_object already fails this check for all ~165 of its other fields, since the merged dict's value type is partly untyped/float; adding two new string fields to the model just grows that existing pile by two. Suppressed at the one call site with a reason, since fixing the root typing is out of scope here. --- .../proxy/guardrails/auto_router_compression.py | 12 ++++++++---- litellm/router.py | 16 +++++++--------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index d26479f7de0..3b0804e40e2 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -13,7 +13,7 @@ each hop sees. """ import contextvars -from collections.abc import Iterable, Mapping, MutableMapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -89,7 +89,7 @@ def policy_for_model( markers: Final = tuple( litellm_params for deployment in deployments - if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) + if isinstance(litellm_params := deployment.get("litellm_params"), Mapping) # pyright: ignore[reportUnnecessaryIsInstance] # filters out non-Mapping and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) ) requested: Final = frozenset(request_tags) @@ -130,7 +130,7 @@ def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: async def arm_pre_call( - data: MutableMapping[str, object], # mutable-ok: arms the live request dict in place + data: dict[str, object], # mutable-ok: arms the live request dict in place llm_router: "Router | None", ) -> None: """Apply an auto router's compression policy, if any, before guardrails run. @@ -183,7 +183,11 @@ async def arm_pre_call( from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - snapshot: Final = resolve_structured_messages(messages=data.get("messages"), request_kwargs=data) + raw_messages: Final = data.get("messages") + snapshot: Final = resolve_structured_messages( + messages=raw_messages if isinstance(raw_messages, list) else None, + request_kwargs=data, + ) if snapshot is not None: _routing_messages_snapshot.set(tuple(MappingProxyType(dict(message)) for message in snapshot)) diff --git a/litellm/router.py b/litellm/router.py index 9637ad98c9a..989914b1610 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8644,7 +8644,7 @@ class Router: raise ValueError(ptu_error) zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None litellm_params: Final[LiteLLM_Params] = LiteLLM_Params( - **( + **( # pyright: ignore[reportArgumentType] # untyped merged dict; already true for every field here _litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing}) @@ -13066,7 +13066,7 @@ class Router: else None ) - pre_routing_hook_response: Final = await selected_strategy.strategy.async_pre_routing_hook( + routed: Final = await selected_strategy.strategy.async_pre_routing_hook( model=registered_model_name, request_kwargs=request_kwargs, messages=routing_messages if routing_messages is not None else messages, @@ -13079,13 +13079,11 @@ class Router: # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, # and pydantic reconstructs a validated list field rather than keeping the # exact object passed in, even when nothing about it changed. - if ( - pre_routing_hook_response is not None - and routing_messages is not None - and pre_routing_hook_response.messages == routing_messages - ): - restored: Final = {"messages": messages} # mutable-ok: pydantic's model_copy takes a dict - pre_routing_hook_response = pre_routing_hook_response.model_copy(update=restored) + pre_routing_hook_response: Final = ( + routed.model_copy(update={"messages": messages}) # mutable-ok: pydantic's model_copy takes a dict + if routed is not None and routing_messages is not None and routed.messages == routing_messages + else routed + ) self._record_routing_decision( request_kwargs=request_kwargs, routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), From 6ee33df952ef9102f14960ed04c46e8f49900d66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:25:12 -0700 Subject: [PATCH 17/61] fix(realtime): relay the upstream websocket close to the client instead of hanging When the provider closes the realtime websocket (for example Vertex Live refusing the session with 1008 "Publisher model ... was not found"), the proxy swallowed the close and kept waiting on the client, so the client sat on an open socket with nothing coming back and the session was logged as a $0 success The backend relay now returns the upstream close, and bidirectional_forward sends the client an OpenAI-style error event naming the upstream code and reason, then closes the client socket with the same code (or 1011 when the upstream code is one a server may not send). A session the upstream refused before sending any frame is logged through the failure handlers instead of as a success --- litellm/litellm_core_utils/realtime_errors.py | 8 + .../litellm_core_utils/realtime_streaming.py | 186 ++++++++++++------ .../test_realtime_errors.py | 10 + .../test_realtime_streaming.py | 169 +++++++++++++++- 4 files changed, 303 insertions(+), 70 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_errors.py b/litellm/litellm_core_utils/realtime_errors.py index e1b957f4325..3c064728a66 100644 --- a/litellm/litellm_core_utils/realtime_errors.py +++ b/litellm/litellm_core_utils/realtime_errors.py @@ -29,3 +29,11 @@ def websocket_close_reason(message: str, fallback: str) -> str: if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES: return message return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + + +def client_close_code(upstream_code: int) -> int: + from websockets.frames import EXTERNAL_CLOSE_CODES, CloseCode + + if upstream_code in EXTERNAL_CLOSE_CODES or 3000 <= upstream_code < 5000: + return upstream_code + return int(CloseCode.INTERNAL_ERROR) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 8479e108d17..746343026ed 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,7 +1,9 @@ import asyncio import json -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +import traceback +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast from typing_extensions import ReadOnly @@ -19,9 +21,11 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging +from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from websockets.exceptions import ConnectionClosed from litellm.types.guardrails import GuardrailEventHooks @@ -30,8 +34,22 @@ else: CLIENT_CONNECTION_CLASS = Any -class _ClientWebSocketExceptions(Protocol): - ConnectionClosed: type[Exception] +@dataclass(frozen=True, slots=True) +class BackendClose: + code: int + reason: str + + @property + def message(self) -> str: + if not self.reason: + return f"upstream websocket closed with code {self.code}" + return f"upstream websocket closed with code {self.code}: {self.reason}" + + +def backend_close_from(error: "ConnectionClosed") -> BackendClose: + if error.rcvd is None: + return BackendClose(code=1006, reason=str(error)) + return BackendClose(code=error.rcvd.code, reason=error.rcvd.reason) class _ASGIScope(TypedDict, total=False): @@ -69,10 +87,13 @@ class _ScopedWebSocket(Protocol): class _ClientWebSocket(_ScopedWebSocket, Protocol): - exceptions: _ClientWebSocketExceptions - async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class _LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... def _decode_json_object(payload: str) -> Mapping[str, object]: @@ -108,11 +129,14 @@ class RealTimeStreaming: backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, + logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj + self._logging_worker = logging_worker self.messages: list[OpenAIRealtimeEvents] = [] + self._backend_sent_frames: bool = False self.input_message: dict = {} self.input_messages: list[dict[str, str]] = [] self.session_tools: list[dict] = [] @@ -388,7 +412,7 @@ class RealTimeStreaming: # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) ) @@ -1035,60 +1059,84 @@ class RealTimeStreaming: return True return False - async def backend_to_client_send_messages(self): + async def _relay_backend_messages(self) -> NoReturn: + while True: + try: + raw_response = await self.backend_ws.recv(decode=False) + except TypeError: + raw_response = await self.backend_ws.recv() + self._backend_sent_frames = True + + if isinstance(raw_response, bytes): + try: + raw_response = raw_response.decode("utf-8") + except UnicodeDecodeError: + verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") + continue + + if self.provider_config: + try: + await self._handle_provider_config_message(raw_response) + except Exception as e: + verbose_logger.exception("Error processing backend message, skipping: %s", e) + continue + else: + event = self._parse_backend_event(raw_response) + if event is None: + await self.websocket.send_text(raw_response) + continue + + if self._should_drop_event_from_client(event): + continue + + if await self._handle_raw_backend_message(event, raw_response): + continue + + event = self._normalize_event_for_ga_client(event) + self.store_message(event) + + if not self._client_wants_beta: + await self.websocket.send_text(json.dumps(event)) + continue + + translated = self._translate_event_to_beta(event) + if translated is None: + continue + await self.websocket.send_text(json.dumps(translated)) + + async def backend_to_client_send_messages(self) -> BackendClose: import websockets try: - while True: - try: - raw_response = await self.backend_ws.recv(decode=False) - except TypeError: - raw_response = await self.backend_ws.recv() - - if isinstance(raw_response, bytes): - try: - raw_response = raw_response.decode("utf-8") - except UnicodeDecodeError: - verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.") - continue - - if self.provider_config: - try: - await self._handle_provider_config_message(raw_response) - except Exception as e: - verbose_logger.exception("Error processing backend message, skipping: %s", e) - continue - else: - event = self._parse_backend_event(raw_response) - if event is None: - await self.websocket.send_text(raw_response) - continue - - if self._should_drop_event_from_client(event): - continue - - if await self._handle_raw_backend_message(event, raw_response): - continue - - event = self._normalize_event_for_ga_client(event) - self.store_message(event) - - if not self._client_wants_beta: - await self.websocket.send_text(json.dumps(event)) - continue - - translated = self._translate_event_to_beta(event) - if translated is None: - continue - await self.websocket.send_text(json.dumps(translated)) - + await self._relay_backend_messages() except websockets.exceptions.ConnectionClosed as e: verbose_logger.exception("Connection closed in backend to client send messages - %s", e) - except Exception as e: - verbose_logger.exception("Error in backend to client send messages: %s", e) - finally: + close: Final = backend_close_from(e) + self._flush_unbilled_transcription_usage() + if self._backend_refused_session(close): + await self.log_backend_refusal(e) + else: + await self.log_messages() + return close + except asyncio.CancelledError: self._flush_unbilled_transcription_usage() await self.log_messages() + raise + except Exception as e: + verbose_logger.exception("Error in backend to client send messages: %s", e) + self._flush_unbilled_transcription_usage() + await self.log_messages() + return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") + + def _backend_refused_session(self, close: BackendClose) -> bool: + return close.code != 1000 and not self._backend_sent_frames and not self.messages + + async def log_backend_refusal(self, error: Exception) -> None: + if not self.logging_obj: + return + self._logging_worker.ensure_initialized_and_enqueue( + self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True) + ) @staticmethod def _detect_beta_header(websocket: _ScopedWebSocket) -> bool: @@ -1484,20 +1532,28 @@ class RealTimeStreaming: except Exception as e: verbose_logger.debug("Error in client ack messages: %s", e) - async def bidirectional_forward(self): + async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages()) + client_task: Final = asyncio.create_task(self.client_ack_messages()) try: - await self.client_ack_messages() - except self.websocket.exceptions.ConnectionClosed: - verbose_logger.debug("Connection closed") - forward_task.cancel() + await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED) + if not client_task.done(): + await self._close_client(forward_task.result()) finally: - if not forward_task.done(): - forward_task.cancel() - try: - await forward_task - except asyncio.CancelledError: - pass + forward_task.cancel() + client_task.cancel() + await asyncio.gather(forward_task, client_task, return_exceptions=True) + + async def _close_client(self, close: BackendClose) -> None: + try: + if close.code != 1000: + await self.websocket.send_text(realtime_error_event(close.message, error_type="server_error")) + await self.websocket.close( + code=client_close_code(close.code), + reason=websocket_close_reason(close.reason, fallback=close.message), + ) + except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way + verbose_logger.debug("Could not relay the upstream close to the client: %s", e) def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool: diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py index 494d16b0b9b..1d2cf905f4e 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py @@ -1,8 +1,10 @@ import json +import pytest from litellm.litellm_core_utils.realtime_errors import ( WEBSOCKET_CLOSE_REASON_MAX_BYTES, + client_close_code, realtime_error_event, websocket_close_reason, ) @@ -42,3 +44,11 @@ def test_websocket_close_reason_truncates_multibyte_message_by_bytes(): assert len(reason.encode("utf-8")) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES assert reason == "あ" * (WEBSOCKET_CLOSE_REASON_MAX_BYTES // 3) assert "�" not in reason + + +@pytest.mark.parametrize( + ("upstream_code", "expected"), + [(1000, 1000), (1008, 1008), (1011, 1011), (4001, 4001), (1005, 1011), (1006, 1011), (1015, 1011), (2999, 1011)], +) +def test_client_close_code_only_forwards_codes_a_server_may_send(upstream_code, expected): + assert client_close_code(upstream_code) == expected diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 52e88db753a..1e0456079e9 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -1,8 +1,13 @@ +import asyncio import json +from collections.abc import Coroutine +from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from websockets.exceptions import ConnectionClosed +from websockets.frames import Close import litellm @@ -2941,13 +2946,11 @@ async def test_log_messages_routes_async_logging_through_bounded_worker(): realtime turn leaves a suspended task pinning its response in memory -> an unbounded leak. Regression for that fix.""" logging_obj = MagicMock() - streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj) + mock_worker = MagicMock() + streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj, logging_worker=mock_worker) streaming.messages = [{"type": "session.created"}] - with ( - patch("litellm.litellm_core_utils.realtime_streaming.GLOBAL_LOGGING_WORKER") as mock_worker, - patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task, - ): + with patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task: await streaming.log_messages() mock_worker.ensure_initialized_and_enqueue.assert_called_once() @@ -3111,3 +3114,159 @@ async def test_session_close_flush_noop_without_unbilled_usage(): isinstance(message, dict) and message.get("type") == "conversation.item.input_audio_transcription.completed" for message in streaming.messages ) + + + +_UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found" + + +class _InlineLoggingWorker: + def __init__(self) -> None: + self.enqueued: tuple[Coroutine[object, object, None], ...] = () + + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: + self.enqueued = (*self.enqueued, async_coroutine) + + async def drain(self) -> None: + for coroutine in self.enqueued: + await coroutine + + +class _RecordingLogging: + def __init__(self) -> None: + self.logged_sessions: tuple[tuple[dict, ...], ...] = () + self.logged_failures: tuple[Exception, ...] = () + + async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None: + self.logged_sessions = (*self.logged_sessions, tuple(result)) + + async def dispatch_failure_handlers( + self, exception: Exception, traceback_exception: str, prefer_async_handlers: bool = False + ) -> None: + self.logged_failures = (*self.logged_failures, exception) + + +@dataclass(frozen=True, slots=True) +class _RelaySession: + streaming: RealTimeStreaming + logging: _RecordingLogging + worker: _InlineLoggingWorker + + async def run(self) -> None: + await asyncio.wait_for(self.streaming.bidirectional_forward(), timeout=2) + await self.worker.drain() + + +async def _wait_forever() -> str: + await asyncio.Event().wait() + raise AssertionError("unreachable") + + +def _client_ws_that_never_sends() -> MagicMock: + client_ws: Final = MagicMock() + client_ws.headers = {} + client_ws.receive_text = AsyncMock(side_effect=_wait_forever) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + return client_ws + + +def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock: + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=list(frames)) + return backend_ws + + +def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession: + logging: Final = _RecordingLogging() + worker: Final = _InlineLoggingWorker() + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker + ) + return _RelaySession(streaming=streaming, logging=logging, worker=worker) + + +def _error_events_sent_to(client_ws: MagicMock) -> list[dict]: + events: Final = (json.loads(call.args[0]) for call in client_ws.send_text.await_args_list) + return [event for event in events if event.get("type") == "error"] + + +@pytest.mark.asyncio +async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): + client_ws: Final = _client_ws_that_never_sends() + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert error_event["error"]["type"] == "server_error" + assert "1008" in error_event["error"]["message"] + assert _UPSTREAM_REFUSAL in error_event["error"]["message"] + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + + +@pytest.mark.asyncio +async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error(): + client_ws: Final = _client_ws_that_never_sends() + session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(None, None))) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert "1006" in error_event["error"]["message"] + client_ws.close.assert_awaited_once() + assert client_ws.close.await_args.kwargs["code"] == 1011 + + +@pytest.mark.asyncio +async def test_bidirectional_forward_relays_normal_upstream_close_without_error_event(): + client_ws: Final = _client_ws_that_never_sends() + session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(Close(1000, ""), None))) + + await session.run() + + assert _error_events_sent_to(client_ws) == [] + client_ws.close.assert_awaited_once() + assert client_ws.close.await_args.kwargs["code"] == 1000 + + +@pytest.mark.asyncio +async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success(): + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + +@pytest.mark.asyncio +async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success(): + client_ws: Final = _client_ws_that_never_sends() + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close)) + + await session.run() + + (logged_session,) = session.logging.logged_sessions + assert [event["type"] for event in logged_session] == ["session.created"] + assert session.logging.logged_failures == () + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + + +@pytest.mark.asyncio +async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close(): + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client went away")) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=_wait_forever) + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + assert session.logging.logged_sessions == ((),) + assert session.logging.logged_failures == () + client_ws.close.assert_not_awaited() From 85d45fbb4b6f9299af75b6891af61664176ad69f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:33:30 -0700 Subject: [PATCH 18/61] fix(realtime): relay the upstream close even when a client message hit the closed socket first When the upstream closes while the proxy is forwarding a client message, the client loop ends before the backend relay sees the close, and the relay skipped closing the client because it read the client loop's exit as the client hanging up. The client loop now reports why it stopped, so a close observed on the backend send still reaches the client with the error event and the upstream close code --- .../litellm_core_utils/realtime_streaming.py | 19 ++++++++-- .../test_realtime_streaming.py | 37 +++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 746343026ed..530391c7b57 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -3,6 +3,7 @@ import json import traceback from collections.abc import Coroutine, Mapping, Sequence from dataclasses import dataclass +from enum import Enum, auto from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast from typing_extensions import ReadOnly @@ -46,6 +47,11 @@ class BackendClose: return f"upstream websocket closed with code {self.code}: {self.reason}" +class ClientLoopExit(Enum): + CLIENT_DISCONNECTED = auto() + BACKEND_CLOSED = auto() + + def backend_close_from(error: "ConnectionClosed") -> BackendClose: if error.rcvd is None: return BackendClose(code=1006, reason=str(error)) @@ -1291,7 +1297,9 @@ class RealTimeStreaming: item["content"] = new_content return item - async def client_ack_messages(self): + async def client_ack_messages(self) -> ClientLoopExit: + import websockets + client_event: _ClientEventFrame try: while True: @@ -1529,16 +1537,21 @@ class RealTimeStreaming: if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True + except websockets.exceptions.ConnectionClosed as e: + verbose_logger.debug("Backend closed while forwarding a client message: %s", e) + return ClientLoopExit.BACKEND_CLOSED except Exception as e: verbose_logger.debug("Error in client ack messages: %s", e) + return ClientLoopExit.CLIENT_DISCONNECTED async def bidirectional_forward(self) -> None: forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages()) client_task: Final = asyncio.create_task(self.client_ack_messages()) try: await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED) - if not client_task.done(): - await self._close_client(forward_task.result()) + if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED: + return + await self._close_client(await forward_task) finally: forward_task.cancel() client_task.cancel() diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 1e0456079e9..41b7557f6b2 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3134,9 +3134,13 @@ class _InlineLoggingWorker: class _RecordingLogging: def __init__(self) -> None: + self.model_call_details: dict[str, object] = {} self.logged_sessions: tuple[tuple[dict, ...], ...] = () self.logged_failures: tuple[Exception, ...] = () + def pre_call(self, input: str | dict, api_key: str) -> None: + return None + async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None: self.logged_sessions = (*self.logged_sessions, tuple(result)) @@ -3257,6 +3261,39 @@ async def test_upstream_close_after_relayed_events_still_logs_the_session_as_suc client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +@pytest.mark.asyncio +async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client(): + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + backend_closed: Final = asyncio.Event() + client_messages: Final = iter((json.dumps({"type": "response.create"}),)) + + async def receive_text() -> str: + message = next(client_messages, None) + return message if message is not None else await _wait_forever() + + async def send_to_backend(_message: str) -> None: + backend_closed.set() + raise upstream_close + + async def recv_from_backend() -> bytes: + await backend_closed.wait() + raise upstream_close + + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = receive_text + backend_ws: Final = MagicMock() + backend_ws.send = send_to_backend + backend_ws.recv = recv_from_backend + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert _UPSTREAM_REFUSAL in error_event["error"]["message"] + client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) + assert session.logging.logged_failures == (upstream_close,) + + @pytest.mark.asyncio async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close(): client_ws: Final = _client_ws_that_never_sends() From c27f1e348dd2f6191177e4b1016388bce1b161f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:47:31 -0700 Subject: [PATCH 19/61] fix(ui): name the object arguments at two new call sites to bring the inline-object lint budget back under its ceiling --- ui/litellm-dashboard/eslint-budgets.json | 2 +- .../src/components/add_model/ClassificationMethodConfig.tsx | 5 +++-- .../add_model/build_complexity_router_config.test.ts | 5 +++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index 3f9163d9028..b3c77e287fc 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -3,7 +3,7 @@ "no-console": { "max": 12, "target": 0 }, "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, - "local/no-large-inline-object-arg": { "max": 555, "target": 300 }, + "local/no-large-inline-object-arg": { "max": 554, "target": 300 }, "local/no-long-condition-chain": { "max": 265, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, "testing-library/no-node-access": { "max": 716, "target": 500 }, diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index cc66103fc86..188ef7f8cb5 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -315,12 +315,13 @@ const ClassificationMethodConfig: React.FC = ({ timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, classification_rubric: selectedRubric, }; - onChange({ + const nextValue: ComplexityRouterConfigValue = { ...value, ...(selectedRubric && { classifier_llm_config: rubricConfig }), classification_prompt: classificationPrompt, classification_examples: classificationExamples, - }); + }; + onChange(nextValue); }; const handleClassifierModelChange = (model: string) => { diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 1b5bb9e72eb..81f05a94a61 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -1170,12 +1170,13 @@ describe("buildComplexityRouterConfig stall escalation", () => { }); it("emits the toggle and both knobs when it is on", () => { - const config = buildComplexityRouterConfig({ + const params: BuildComplexityRouterConfigParams = { ...baseParams, stallEscalationEnabled: true, stallEscalationWindow: 8, stallEscalationRepeatThreshold: 4, - }); + }; + const config = buildComplexityRouterConfig(params); expect(config.stall_escalation_enabled).toBe(true); expect(config.stall_escalation_window).toBe(8); expect(config.stall_escalation_repeat_threshold).toBe(4); From da9dbdba961ce2981f9d82669ab9e3eb3a9d90a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 19:52:48 -0700 Subject: [PATCH 20/61] fix(realtime): treat any client receive failure as a client hangup client_ack_messages classified a websockets ConnectionClosed raised by the client socket as the backend closing, so bidirectional_forward kept waiting on the upstream instead of ending the session. Starlette clients raise WebSocketDisconnect, but the realtime test client in tests/llm_translation/realtime raises websockets.exceptions.ConnectionClosed, which hung test_openai_realtime_simple.py until the run was killed. Only the receive_text call now maps every exception to CLIENT_DISCONNECTED; the loop body keeps ConnectionClosed as BACKEND_CLOSED, since the backend socket is the only websockets socket touched there. --- litellm/litellm_core_utils/realtime_streaming.py | 11 ++++++++++- .../litellm_core_utils/test_realtime_streaming.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 530391c7b57..bb7fbd81146 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1297,13 +1297,22 @@ class RealTimeStreaming: item["content"] = new_content return item + async def _receive_client_message(self) -> str | None: + try: + return await self.websocket.receive_text() + except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone + verbose_logger.debug("Client disconnected: %s", e) + return None + async def client_ack_messages(self) -> ClientLoopExit: import websockets client_event: _ClientEventFrame try: while True: - message = await self.websocket.receive_text() + message = await self._receive_client_message() + if message is None: + return ClientLoopExit.CLIENT_DISCONNECTED ## GUARDRAIL: intercept conversation.item.create for text-based injection. guardrail_turn_detection_injected = False diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 41b7557f6b2..00addb613c2 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3307,3 +3307,18 @@ async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close( assert session.logging.logged_sessions == ((),) assert session.logging.logged_failures == () client_ws.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the_backend_closing(): + client_ws: Final = _client_ws_that_never_sends() + client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(None, None)) + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=_wait_forever) + session: Final = _relay_session(client_ws, backend_ws) + + await session.run() + + assert session.logging.logged_sessions == ((),) + assert session.logging.logged_failures == () + client_ws.close.assert_not_awaited() From 6385c7b3c53617fb480f5c8875a9b45538174ddf Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 4 Sep 2026 20:28:17 -0700 Subject: [PATCH 21/61] fix(auto-router compression): close three review findings on the per-hop policy Suppression state moves out of request metadata into a request-scoped ContextVar. refresh_proxy_server_request_body_snapshot copies metadata into proxy_server_request.body, which deployments persist to spend logs, so the marker naming each suppressed guardrail was readable by the caller whose request produced it. Recovering it was enough to replay {token}:{name} for any CustomGuardrail and switch off a PII or content-filter guardrail, since the check never verified the named guardrail was a compression one. Nothing is read from metadata now, so there is no marker to forge and the per-process token is no longer needed. Routing-side compression reads the live messages instead of a pre-guardrail copy. arm_pre_call runs before the pre-call hook, so its snapshot held the prompt as it was before any masking guardrail rewrote it, and messages_for_routing handed that to a compression guardrail which POSTs it to an external service. Masked content left the proxy anyway. The cost is one combination: when the model hop compressed and the hops differ, routing now classifies on the compressed text, since no uncompressed copy survives that a masking guardrail has already seen. policy_for_model no longer falls back to a marker scoped to tags the request does not carry, which applied an 'eu' policy to a 'us' request on config order alone. Each fix carries a regression test; all three fail when the fix is reverted. --- litellm/constants.py | 1 - litellm/integrations/custom_guardrail.py | 36 +- .../guardrails/auto_router_compression.py | 101 +++--- .../integrations/test_custom_guardrail.py | 335 +++++------------- .../test_auto_router_compression.py | 147 ++++---- tests/test_litellm/test_router.py | 25 +- 6 files changed, 236 insertions(+), 409 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 25fdaec20de..43fefaae048 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -217,7 +217,6 @@ PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" # Metadata key listing compression guardrails an auto router's own compression # policy suppresses for this request. See litellm.proxy.guardrails.auto_router_compression. -AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: Final = "_auto_router_suppressed_compression_guardrails" # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1ec08641706..f511d128dfc 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -45,7 +45,6 @@ dc: Final = DualCache() from litellm.constants import ( - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY, GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) @@ -941,33 +940,22 @@ class CustomGuardrail(CustomLogger): """ return False - def auto_router_suppression_marker(self) -> str | None: - """The value `arm_pre_call` must write to suppress this guardrail. + def _suppressed_by_auto_router_compression(self) -> bool: + """True when an auto router's own compression policy suppresses this guardrail. - Carries the per-process token for the same reason `_pre_call_marker` does: a - caller controls request metadata, so a bare guardrail name there would let any - request switch off a PII, content-filter, or compression guardrail for itself. - The token is never sent to the caller, so the marker cannot be forged. + Reads request-scoped state set by `arm_pre_call`, never request metadata. The + caller controls metadata, and metadata reaches spend logs the caller can read, + so a suppression list carried there would be one a request could replay to + switch off a PII or content-filter guardrail for itself. """ name: Final = self.guardrail_name if not name: - return None - return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - - def _suppressed_by_auto_router_compression(self, data: Mapping[str, object]) -> bool: - """True when an auto router's own compression policy suppresses this guardrail.""" - marker: Final = self.auto_router_suppression_marker() - if marker is None: return False - for meta_key in ("metadata", "litellm_metadata"): - meta = data.get(meta_key) - if isinstance(meta, Mapping): - # arm_pre_call writes a tuple; it arrives as a list once the metadata - # has been round-tripped through JSON. - suppressed = meta.get(AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY) - if isinstance(suppressed, (list, tuple)) and marker in suppressed: - return True - return False + from litellm.proxy.guardrails.auto_router_compression import ( + suppressed_compression_guardrails, + ) + + return name in suppressed_compression_guardrails() def should_run_guardrail( self, @@ -977,7 +965,7 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ - if self._suppressed_by_auto_router_compression(data): + if self._suppressed_by_auto_router_compression(): return False requested_guardrails: Final = self.get_guardrail_from_metadata(data) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 3b0804e40e2..2a4a0c82022 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -15,11 +15,9 @@ each hop sees. import contextvars from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass -from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs @@ -31,16 +29,22 @@ if TYPE_CHECKING: COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# The pre-compression messages, so a routing decision that does not share the model -# call's compression still classifies on the original text. Deliberately a ContextVar -# rather than a metadata key: `refresh_proxy_server_request_body_snapshot` copies -# metadata into `proxy_server_request.body`, which deployments persist to spend logs, -# and this holds the prompt as it was before any masking guardrail rewrote it. -_routing_messages_snapshot: Final[contextvars.ContextVar[tuple[Mapping[str, object], ...] | None]] = ( - contextvars.ContextVar("litellm_auto_router_routing_messages_snapshot", default=None) +# Compression guardrails this request's auto router has switched off. Deliberately a +# ContextVar rather than a metadata key: `refresh_proxy_server_request_body_snapshot` +# copies metadata into `proxy_server_request.body`, which deployments persist to spend +# logs. A suppression list that reaches a log the caller can read is a list the caller +# can replay, which would let any request switch off a PII or content-filter guardrail. +# Nothing here is caller-supplied, so there is no marker to forge in the first place. +_suppressed_compression_guardrails: Final[contextvars.ContextVar[frozenset[str]]] = contextvars.ContextVar( + "litellm_auto_router_suppressed_compression_guardrails", default=frozenset() ) +def suppressed_compression_guardrails() -> frozenset[str]: + """Names of the compression guardrails this request's auto router suppresses.""" + return _suppressed_compression_guardrails.get() + + @dataclass(frozen=True, slots=True) class AutoRouterCompressionPolicy: """An auto router's compression choice for each hop. ``None`` means no compression.""" @@ -96,7 +100,11 @@ def policy_for_model( tag_matched: Final = tuple( params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) - for params in (*tag_matched, *markers): + # Only untagged markers may serve as the fallback. A marker scoped to tags this + # request does not carry describes a different slice of traffic, so falling back + # to it would apply, say, an "eu" policy to a "us" request purely on config order. + untagged: Final = tuple(params for params in markers if not params.get("tags")) + for params in (*tag_matched, *untagged): policy = policy_from_litellm_params(params) if policy is not None: return policy @@ -135,12 +143,10 @@ async def arm_pre_call( ) -> None: """Apply an auto router's compression policy, if any, before guardrails run. - Suppresses every other compression guardrail, re-enables the model-side - guardrail the policy names (if any) even when it isn't ``default_on``, and - snapshots the pre-compression messages so the routing decision can read them - independently of whatever the model-side guardrail does to `data`. + Suppresses every other compression guardrail and re-enables the model-side + guardrail the policy names (if any) even when it isn't ``default_on``. """ - _routing_messages_snapshot.set(None) + _suppressed_compression_guardrails.set(frozenset()) if llm_router is None: return @@ -162,18 +168,16 @@ async def arm_pre_call( if policy is None: return - _, metadata = get_or_create_metadata_bucket(data) - # Markers carry a per-process token so a caller cannot suppress a guardrail by - # naming it in its own request metadata. - suppressed: Final = tuple( - marker - for guardrail in _active_compression_guardrails() - if guardrail.guardrail_name != policy.model and (marker := guardrail.auto_router_suppression_marker()) + _suppressed_compression_guardrails.set( + frozenset( + name + for guardrail in _active_compression_guardrails() + if (name := guardrail.guardrail_name) and name != policy.model + ) ) - if suppressed: - metadata[AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] = suppressed if policy.model is not None: + _, metadata = get_or_create_metadata_bucket(data) requested: Final = metadata.get("guardrails") existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () if policy.model not in existing: @@ -181,20 +185,6 @@ async def arm_pre_call( # isinstance(..., list) and extends it, and would drop a tuple on the floor. metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list - from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - - raw_messages: Final = data.get("messages") - snapshot: Final = resolve_structured_messages( - messages=raw_messages if isinstance(raw_messages, list) else None, - request_kwargs=data, - ) - if snapshot is not None: - _routing_messages_snapshot.set(tuple(MappingProxyType(dict(message)) for message in snapshot)) - - -def _snapshot_messages() -> tuple[Mapping[str, object], ...] | None: - return _routing_messages_snapshot.get() - def _as_routing_messages( messages: Iterable[Mapping[str, object]], @@ -213,23 +203,22 @@ async def messages_for_routing( """Messages to use for a routing decision, per `policy.routing`. Returns None when the caller should route on whatever messages it already has. - The model call is untouched either way: model-side compression, if any, already - ran as an ordinary pre-call guardrail before the router was reached, so when the - two hops differ the routing decision reads the pre-compression snapshot rather - than what that guardrail left behind. + + Always reads the live messages, never a pre-guardrail copy of them. The routing + hop compresses through a real guardrail, which POSTs the text to an external + compression service, so it must see what every other guardrail has already done + to the request. Routing on a snapshot taken before the pre-call hook would send + a masking guardrail's own input straight back out of the proxy. + + The consequence, when the model hop compressed and the two hops differ: the + messages in hand are that guardrail's output, and there is no un-compressed copy + left to route on. The routing decision reads the compressed text in that one + combination rather than leaking the original. """ - if policy is None: + if policy is None or policy.routing is None: return None - snapshot: Final = _snapshot_messages() - original: Final = snapshot if snapshot is not None else messages - - if policy.routing is None: - # Explicitly no compression for routing. When the model side compressed, the - # messages in hand are its output, so fall back to the untouched snapshot. - return _as_routing_messages(snapshot) if policy.model is not None and snapshot is not None else None - - if not original: + if not messages: return None from litellm.proxy.common_utils.registry_read_through import ( @@ -241,20 +230,20 @@ async def messages_for_routing( verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' not found; routing on uncompressed messages", policy.routing ) - return _as_routing_messages(original) + return _as_routing_messages(messages) inputs: Final[GenericGuardrailAPIInputs] = { - "structured_messages": _as_routing_messages(original) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape + "structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape } model: Final = request_kwargs.get("model") # A throwaway request_data: apply_guardrail writes its stats onto this dict, not the # real request's metadata, so routing-side compression never double-counts against # extract_compression_saved_tokens's model-savings accounting. - stats_sink: Final = {"messages": original, "model": model} # mutable-ok: apply_guardrail writes its stats here + stats_sink: Final = {"messages": messages, "model": model} # mutable-ok: apply_guardrail writes its stats here result: Final = await guardrail.apply_guardrail( inputs=inputs, request_data=stats_sink, input_type="request", ) compressed: Final = result.get("structured_messages") - return compressed if isinstance(compressed, list) else _as_routing_messages(original) + return compressed if isinstance(compressed, list) else _as_routing_messages(messages) diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index f590903cb74..c4a702d453c 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -13,7 +13,6 @@ from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetai class TestCustomGuardrailDeploymentHook: - @pytest.mark.asyncio async def test_async_pre_call_deployment_hook_no_guardrails(self): """Test that method returns kwargs unchanged when no guardrails are present""" @@ -26,18 +25,14 @@ class TestCustomGuardrailDeploymentHook: "guardrails": None, } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs # Test with guardrails as non-list kwargs["guardrails"] = "not_a_list" - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert result == kwargs @@ -64,9 +59,7 @@ class TestCustomGuardrailDeploymentHook: "user_api_key_request_route": "test_route", } - result = await custom_guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + result = await custom_guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) # Verify async_pre_call_hook was called with correct parameters custom_guardrail.async_pre_call_hook.assert_called_once() @@ -99,9 +92,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -114,9 +105,7 @@ class TestCustomGuardrailDeploymentHook: } guardrail.mark_pre_call_hook_ran(kwargs) - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 0 @@ -130,9 +119,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -144,9 +131,7 @@ class TestCustomGuardrailDeploymentHook: "metadata": {}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 @@ -175,9 +160,7 @@ class TestCustomGuardrailDeploymentHook: super().__init__(guardrail_name="g1", default_on=True) self.pre_call_count = 0 - async def async_pre_call_hook( - self, user_api_key_dict, cache, data, call_type - ): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.pre_call_count += 1 return data @@ -189,15 +172,12 @@ class TestCustomGuardrailDeploymentHook: "metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]}, } - await guardrail.async_pre_call_deployment_hook( - kwargs=kwargs, call_type=CallTypes.completion - ) + await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.completion) assert guardrail.pre_call_count == 1 class TestCustomGuardrailShouldRunGuardrail: - def test_should_run_guardrail_with_litellm_metadata(self): """Test that should_run_guardrail works with litellm_metadata pattern""" from litellm.types.guardrails import GuardrailEventHooks @@ -214,9 +194,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -236,9 +214,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"guardrails": ["test_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -255,9 +231,7 @@ class TestCustomGuardrailShouldRunGuardrail: # Test with guardrails at root level data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]} - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True @@ -277,9 +251,7 @@ class TestCustomGuardrailShouldRunGuardrail: "litellm_metadata": {"guardrails": ["different_guardrail"]}, } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is False @@ -298,9 +270,7 @@ class TestCustomGuardrailShouldRunGuardrail: "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "test"}], } - result = custom_guardrail.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) + result = custom_guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) assert result is True, "Global guardrail should run when default_on=True" # Test 2: User-injected disable at root level is IGNORED @@ -312,9 +282,7 @@ class TestCustomGuardrailShouldRunGuardrail: result = custom_guardrail.should_run_guardrail( data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call ) - assert ( - result is True - ), "User-injected disable_global_guardrails should be ignored" + assert result is True, "User-injected disable_global_guardrails should be ignored" # Test 3: User-injected disable in metadata is IGNORED data_with_disable_metadata = { @@ -345,12 +313,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {"request_tags": ["user-supplied"]}, } - result = custom_guardrail.should_run_guardrail( - data=data_cross_key, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" + result = custom_guardrail.should_run_guardrail(data=data_cross_key, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must not be shadowed by user-supplied litellm_metadata" # Test 6: After the pre-call strip runs, user-injected # user_api_key_metadata in the non-authoritative metadata key is gone. @@ -361,12 +325,8 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, "litellm_metadata": {}, # post-strip: attacker payload removed } - result = custom_guardrail.should_run_guardrail( - data=data_post_strip, event_type=GuardrailEventHooks.pre_call - ) - assert ( - result is False - ), "Admin config in metadata must be respected when other metadata key is empty" + result = custom_guardrail.should_run_guardrail(data=data_post_strip, event_type=GuardrailEventHooks.pre_call) + assert result is False, "Admin config in metadata must be respected when other metadata key is empty" def test_should_run_guardrail_key_disable_global_not_overruled_by_team_guardrail_list( self, @@ -432,12 +392,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "opted_out_global_guardrails": ["global_guardrail"], } - assert ( - custom_guardrail.should_run_guardrail( - data=data_root, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_root, event_type=GuardrailEventHooks.pre_call) is True # Test 2: User-injected opt-out in metadata is IGNORED data_metadata = { @@ -446,10 +401,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": ["global_guardrail"]}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_metadata, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_metadata, event_type=GuardrailEventHooks.pre_call) is True ) # Test 4: a different guardrail in the opt-out list → still runs @@ -458,12 +410,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_other, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_other, event_type=GuardrailEventHooks.pre_call) is True # Test 5: empty opt-out list → still runs data_empty = { @@ -471,12 +418,7 @@ class TestCustomGuardrailShouldRunGuardrail: "messages": [{"role": "user", "content": "test"}], "metadata": {"opted_out_global_guardrails": []}, } - assert ( - custom_guardrail.should_run_guardrail( - data=data_empty, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert custom_guardrail.should_run_guardrail(data=data_empty, event_type=GuardrailEventHooks.pre_call) is True # Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs data_malformed = { @@ -485,10 +427,7 @@ class TestCustomGuardrailShouldRunGuardrail: "metadata": {"opted_out_global_guardrails": True}, } assert ( - custom_guardrail.should_run_guardrail( - data=data_malformed, event_type=GuardrailEventHooks.pre_call - ) - is True + custom_guardrail.should_run_guardrail(data=data_malformed, event_type=GuardrailEventHooks.pre_call) is True ) def test_should_run_guardrail_opt_out_does_not_affect_non_global(self): @@ -511,17 +450,12 @@ class TestCustomGuardrailShouldRunGuardrail: "guardrails": ["opt_in_guardrail"], }, } - assert ( - non_global.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.pre_call - ) - is True - ) + assert non_global.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is True def test_should_run_guardrail_suppressed_by_auto_router_compression(self): """An auto router's own compression policy can suppress an otherwise-eligible guardrail, even one that is default_on and explicitly requested.""" - from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + from litellm.proxy.guardrails import auto_router_compression from litellm.types.guardrails import GuardrailEventHooks always_on = CustomGuardrail( @@ -529,22 +463,17 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) - data = { - "model": "smart-router", - "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ - always_on.auto_router_suppression_marker() - ], - }, - } + token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"headroom-default"})) + try: + assert ( + always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call) + is False + ) + finally: + auto_router_compression._suppressed_compression_guardrails.reset(token) - assert ( - always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) - is False - ) - - def test_should_run_guardrail_suppression_list_does_not_affect_other_names(self): - from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + def test_should_run_guardrail_suppression_does_not_affect_other_names(self): + from litellm.proxy.guardrails import auto_router_compression from litellm.types.guardrails import GuardrailEventHooks always_on = CustomGuardrail( @@ -552,25 +481,20 @@ class TestCustomGuardrailShouldRunGuardrail: default_on=True, event_hook=GuardrailEventHooks.pre_call, ) - other = CustomGuardrail(guardrail_name="some-other-guardrail", default_on=True) - data = { - "model": "smart-router", - "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ - other.auto_router_suppression_marker() - ], - }, - } + token = auto_router_compression._suppressed_compression_guardrails.set(frozenset({"some-other-guardrail"})) + try: + assert ( + always_on.should_run_guardrail(data={"model": "smart-router"}, event_type=GuardrailEventHooks.pre_call) + is True + ) + finally: + auto_router_compression._suppressed_compression_guardrails.reset(token) - assert ( - always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) - is True - ) - - def test_should_run_guardrail_ignores_a_forged_suppression_marker(self): - """A caller controls request metadata, so a bare guardrail name there must not - switch off an always-on guardrail: only the per-process marker counts.""" - from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY + def test_request_metadata_can_never_suppress_a_guardrail(self): + """Regression (security): suppression state is request-scoped and server-set, + never read from metadata. Metadata reaches spend logs the caller can read, so + anything honored from there is something a later request could replay to switch + off a PII or content-filter guardrail for itself.""" from litellm.types.guardrails import GuardrailEventHooks always_on = CustomGuardrail( @@ -581,17 +505,14 @@ class TestCustomGuardrailShouldRunGuardrail: forged = { "model": "smart-router", "metadata": { - AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: [ + "_auto_router_suppressed_compression_guardrails": [ "headroom-default", - "forged-token:headroom-default", + "any-token:headroom-default", ], }, } - assert ( - always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) - is True - ) + assert always_on.should_run_guardrail(data=forged, event_type=GuardrailEventHooks.pre_call) is True class TestApplyGuardrailCheck: @@ -630,35 +551,33 @@ class TestApplyGuardrailCheck: child_with_override = ChildGuardrailWithOverride() # Test: CustomGuardrail itself has apply_guardrail in its __dict__ - assert ( - "apply_guardrail" in type(CustomGuardrail()).__dict__ - ), "CustomGuardrail should have apply_guardrail in its own __dict__" + assert "apply_guardrail" in type(CustomGuardrail()).__dict__, ( + "CustomGuardrail should have apply_guardrail in its own __dict__" + ) # Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(parent_instance).__dict__ - ), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(parent_instance).__dict__, ( + "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__ - assert ( - "apply_guardrail" not in type(child_without_override).__dict__ - ), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + assert "apply_guardrail" not in type(child_without_override).__dict__, ( + "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)" + ) # Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__ - assert ( - "apply_guardrail" in type(child_with_override).__dict__ - ), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + assert "apply_guardrail" in type(child_with_override).__dict__, ( + "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)" + ) # Verify that all instances still have the method via inheritance (hasattr) - assert hasattr( - parent_instance, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_without_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" - assert hasattr( - child_with_override, "apply_guardrail" - ), "All instances should have apply_guardrail via inheritance" + assert hasattr(parent_instance, "apply_guardrail"), "All instances should have apply_guardrail via inheritance" + assert hasattr(child_without_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) + assert hasattr(child_with_override, "apply_guardrail"), ( + "All instances should have apply_guardrail via inheritance" + ) class TestGuardrailLoggingAggregation: @@ -685,11 +604,7 @@ class TestGuardrailLoggingAggregation: def test_appends_to_existing_metadata_list(self): request_data = { - "metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "existing_guardrail"} - ] - } + "metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "existing_guardrail"}]} } self._invoke_add_log(request_data) @@ -701,11 +616,7 @@ class TestGuardrailLoggingAggregation: assert info[1]["guardrail_name"] == "test_guardrail" def test_converts_existing_metadata_dict_to_list(self): - request_data = { - "metadata": { - "standard_logging_guardrail_information": {"guardrail_name": "legacy"} - } - } + request_data = {"metadata": {"standard_logging_guardrail_information": {"guardrail_name": "legacy"}}} self._invoke_add_log(request_data) @@ -717,18 +628,12 @@ class TestGuardrailLoggingAggregation: def test_appends_to_litellm_metadata(self): request_data = { - "litellm_metadata": { - "standard_logging_guardrail_information": [ - {"guardrail_name": "litellm_existing"} - ] - } + "litellm_metadata": {"standard_logging_guardrail_information": [{"guardrail_name": "litellm_existing"}]} } self._invoke_add_log(request_data) - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert isinstance(info, list) assert len(info) == 2 assert info[1]["guardrail_name"] == "test_guardrail" @@ -745,12 +650,10 @@ class TestGuardrailLoggingAggregation: self._invoke_add_log(request_data) - assert ( - "standard_logging_guardrail_information" not in request_data["metadata"] - ), "entry landed in the caller's metadata, where the spend log does not read it" - info = request_data["litellm_metadata"][ - "standard_logging_guardrail_information" - ] + assert "standard_logging_guardrail_information" not in request_data["metadata"], ( + "entry landed in the caller's metadata, where the spend log does not read it" + ) + info = request_data["litellm_metadata"]["standard_logging_guardrail_information"] assert len(info) == 1 assert info[0]["guardrail_name"] == "test_guardrail" @@ -768,9 +671,7 @@ class TestGuardrailLoggingAggregation: } self._invoke_add_log(request_data) - add_guardrail_to_applied_guardrails_header( - request_data=request_data, guardrail_name="test_guardrail" - ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name="test_guardrail") buckets = { key @@ -816,9 +717,7 @@ class TestGuardrailOtelSpanEmission: assert len(captured) == 1 emitted = captured[0] - recorded = request_data["metadata"]["standard_logging_guardrail_information"][ - -1 - ] + recorded = request_data["metadata"]["standard_logging_guardrail_information"][-1] assert emitted is recorded assert emitted["guardrail_name"] == "emit_guard" assert emitted["start_time"] == 1.0 @@ -828,9 +727,7 @@ class TestGuardrailOtelSpanEmission: def _boom(_entry): raise RuntimeError("otel exporter down") - monkeypatch.setattr( - "litellm.integrations.otel.logger.emit_guardrail_span", _boom - ) + monkeypatch.setattr("litellm.integrations.otel.logger.emit_guardrail_span", _boom) request_data = {"metadata": {}} self._record(self._make_guardrail(), request_data) @@ -927,9 +824,7 @@ class TestGuardrailSensitiveFieldStripping: duration=1.0, ) - logged_response = request_data["metadata"][ - "standard_logging_guardrail_information" - ][0]["guardrail_response"] + logged_response = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert "secret_fields" not in logged_response assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response) @@ -942,9 +837,7 @@ class TestGuardrailSensitiveFieldStripping: guardrail_json_response=[ { "result": "ok", - "secret_fields": { - "raw_headers": {"authorization": "Bearer sk-secret"} - }, + "secret_fields": {"raw_headers": {"authorization": "Bearer sk-secret"}}, }, {"result": "also_ok"}, ], @@ -998,9 +891,7 @@ class TestGuardrailResponseCredentialMasking: duration=1.0, ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] masked_key = logged["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] assert masked_key != plaintext_key @@ -1009,10 +900,7 @@ class TestGuardrailResponseCredentialMasking: assert logged["model"] == "gpt-4o-mini" assert logged["messages"] == [{"role": "user", "content": "hi"}] - assert ( - logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] - == "proj-name" - ) + assert logged["metadata_snapshot"]["callback_vars"]["langsmith_project"] == "proj-name" def test_nested_user_api_key_auth_metadata_is_masked(self): import json @@ -1071,9 +959,7 @@ class TestGuardrailResponseCredentialMasking: request_data: dict = {"metadata": {}} guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}] - }, + guardrail_json_response={"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}, request_data=request_data, guardrail_status="success", ) @@ -1096,9 +982,7 @@ class TestGuardrailResponseCredentialMasking: guardrail_status="success", ) - logged = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ] + logged = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"] assert logged["flagged"] is True assert logged["score"] == 0.94 assert logged["tokens_used"] == 42 @@ -1110,18 +994,14 @@ class TestGuardrailResponseCredentialMasking: plaintext = "lsv2_pt_abcdef1234567890" guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={ - "metadata_snapshot": { - "callback_vars": {"langsmith_api_key": plaintext} - } - }, + guardrail_json_response={"metadata_snapshot": {"callback_vars": {"langsmith_api_key": plaintext}}}, request_data=request_data, guardrail_status="success", ) - masked = request_data["metadata"]["standard_logging_guardrail_information"][0][ - "guardrail_response" - ]["metadata_snapshot"]["callback_vars"]["langsmith_api_key"] + masked = request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_response"][ + "metadata_snapshot" + ]["callback_vars"]["langsmith_api_key"] assert masked != plaintext assert masked.startswith(plaintext[:4]) assert masked.endswith(plaintext[-4:]) @@ -1615,9 +1495,7 @@ class TestEventTypeLogging: guardrail = TestGuardrail() request_data = {"metadata": {}} - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1, ( @@ -1659,9 +1537,7 @@ class TestEventTypeLogging: request_data = {"metadata": {}} with pytest.raises(ValueError, match="blocked"): - await guardrail.apply_guardrail( - inputs={"texts": ["x"]}, request_data=request_data - ) + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data) logged_info = request_data["metadata"]["standard_logging_guardrail_information"] assert len(logged_info) == 1 @@ -1790,9 +1666,7 @@ class TestTracingFieldsPopulation: guardrail_json_response="blocked", request_data=request_data, guardrail_status="guardrail_intervened", - tracing_detail=GuardrailTracingDetail( - policy_template="EU AI Act Article 5" - ), + tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"), ) slg_list = request_data["metadata"]["standard_logging_guardrail_information"] @@ -1834,13 +1708,7 @@ class TestCustomGuardrailSpendLogMatchRedaction: cg = CustomGuardrail(guardrail_name="test-rail") raw = { "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ] } request_data: dict = {"metadata": {}} @@ -1851,17 +1719,10 @@ class TestCustomGuardrailSpendLogMatchRedaction: ) slg = request_data["metadata"]["standard_logging_guardrail_information"][0] assert ( - slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) - assert ( - raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "GG" - ) + assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "GG" def test_add_standard_logging_redacts_regex_field(self): cg = CustomGuardrail(guardrail_name="test-rail") diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index de667e8ed48..0b47e56cb02 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -4,15 +4,16 @@ Unit tests for litellm.proxy.guardrails.auto_router_compression. Covers: - policy_from_litellm_params: absent keys mean no policy; the "none" sentinel normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias, and - picks the tag-scoped marker the request's tags actually match +- policy_for_model: finds the auto-router marker deployment for an alias, picks the + tag-scoped marker the request's tags actually match, and never falls back to a + marker scoped to tags the request does not carry - arm_pre_call: no-op without a policy; suppresses active compression guardrails - with a forgery-proof marker; arms the model-side guardrail even when it isn't - default_on; keeps the pre-compression snapshot out of persisted metadata -- messages_for_routing: no-op without a policy; routes on the pre-compression - snapshot when the two hops differ; compresses via the named guardrail's - apply_guardrail; never writes stats onto the caller's own request_kwargs - (regression for double-counted compression savings) + through request-scoped state rather than metadata, which reaches spend logs a + caller can read; arms the model-side guardrail even when it isn't default_on +- messages_for_routing: no-op without a policy; compresses the live messages every + earlier guardrail has already rewritten, never a pre-guardrail copy of them; + never writes stats onto the caller's own request_kwargs (regression for + double-counted compression savings) """ import json @@ -20,7 +21,6 @@ from typing import Any import pytest -from litellm.constants import AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails import auto_router_compression from litellm.proxy.guardrails.auto_router_compression import ( @@ -136,6 +136,26 @@ class TestPolicyForModel: ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) + def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self): + """Regression: an "eu" marker describes a different slice of traffic, so a "us" + request must not fall back to its policy just because it is configured first.""" + router = _FakeRouter( + [ + _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), + _marker({"auto_router_routing_compression": "headroom-default"}), + ] + ) + policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) + + def test_no_untagged_fallback_means_no_policy(self): + """With only tag-scoped markers and none matching, there is no policy to apply: + inheriting an unrelated slice's compression is worse than inheriting nothing.""" + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) + assert ( + policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + ) + def test_tag_scoped_marker_takes_precedence_over_untagged(self): """Regression: when multiple markers exist, the tag-scoped one the request actually matches should be used, not the first untagged one.""" @@ -220,25 +240,30 @@ class TestArmPreCall: ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} await arm_pre_call(data=data, llm_router=router) - suppressed = data["metadata"][AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY] - assert tuple(suppressed) == (always_on.auto_router_suppression_marker(),) - # The bare name alone must never suppress: that is what a caller could forge. - assert "always-on-compression" not in suppressed + assert auto_router_compression.suppressed_compression_guardrails() == frozenset({"always-on-compression"}) + # Suppression state must never ride along in metadata: that reaches spend + # logs the caller can read, and anything there is replayable. + assert "always-on-compression" not in json.dumps(data.get("metadata", {})) assert always_on.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is False finally: litellm.logging_callback_manager.remove_callback_from_all_lists(always_on) @pytest.mark.asyncio - async def test_a_caller_cannot_suppress_a_guardrail_by_naming_it_in_metadata(self): - """Regression: request metadata is caller-controlled, so a bare guardrail name - there must not switch off a PII, content-filter, or compression guardrail.""" + async def test_suppression_state_never_enters_request_metadata(self): + """Regression (security): a suppression list written to metadata is copied into + proxy_server_request.body and persisted to spend logs, so a caller could read it + back and replay it to switch off a PII or content-filter guardrail.""" guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") - forged = { - "model": "smart-router", - "metadata": {AUTO_ROUTER_SUPPRESSED_COMPRESSION_GUARDRAILS_KEY: ["always-on-compression"]}, - } + import litellm - assert guardrail._suppressed_by_auto_router_compression(forged) is False + litellm.logging_callback_manager.add_litellm_callback(guardrail) + try: + router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) + data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} + await arm_pre_call(data=data, llm_router=router) + assert "suppress" not in json.dumps(data).lower() + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) @pytest.mark.asyncio async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): @@ -259,52 +284,20 @@ class TestArmPreCall: assert data["metadata"]["guardrails"] == ["headroom-b"] @pytest.mark.asyncio - async def test_snapshot_never_lands_in_persisted_metadata(self): - """Regression: refresh_proxy_server_request_body_snapshot copies metadata into - proxy_server_request.body, which deployments persist to spend logs. The - pre-compression snapshot holds the prompt before any masking guardrail ran, so - it must live outside anything that gets serialized.""" + async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): + """Regression (security): arm_pre_call runs before the pre-call guardrails, so + any copy of the messages it retained would be the pre-masking text. Routing-side + compression POSTs its input to an external service, so that copy must not exist.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - original_messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] - data = {"model": "smart-router", "messages": original_messages} + data = {"model": "smart-router", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} await arm_pre_call(data=data, llm_router=router) - assert "123-45-6789" not in json.dumps(data["metadata"]) - assert [dict(m) for m in auto_router_compression._snapshot_messages()] == original_messages - - @pytest.mark.asyncio - async def test_snapshot_is_a_copy_not_the_live_message_list(self): - router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - original_messages = [{"role": "user", "content": "hi"}] - - await arm_pre_call(data={"model": "smart-router", "messages": original_messages}, llm_router=router) - original_messages[0]["content"] = "mutated after the snapshot" - - assert [dict(m) for m in auto_router_compression._snapshot_messages()] == [{"role": "user", "content": "hi"}] - - @pytest.mark.asyncio - async def test_a_request_without_a_policy_clears_a_previous_snapshot(self): - router_with = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) - await arm_pre_call( - data={"model": "smart-router", "messages": [{"role": "user", "content": "first"}]}, llm_router=router_with - ) - - router_without = _FakeRouter([{"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - await arm_pre_call( - data={"model": "plain", "messages": [{"role": "user", "content": "second"}]}, llm_router=router_without - ) - - assert auto_router_compression._snapshot_messages() is None + assert "123-45-6789" not in json.dumps(data.get("metadata", {})) + assert not hasattr(auto_router_compression, "_routing_messages_snapshot") class TestMessagesForRouting: - @pytest.fixture(autouse=True) - def _clear_snapshot(self): - auto_router_compression._routing_messages_snapshot.set(None) - yield - auto_router_compression._routing_messages_snapshot.set(None) - @pytest.mark.asyncio async def test_no_policy_returns_none(self): assert await messages_for_routing(policy=None, messages=[], request_kwargs={}) is None @@ -316,18 +309,15 @@ class TestMessagesForRouting: assert await messages_for_routing(policy=policy, messages=[], request_kwargs={}) is None @pytest.mark.asyncio - async def test_routing_none_with_model_compression_routes_on_the_snapshot(self): - """Regression: with routing explicitly off and the model side compressed, the - messages in hand are the model-side guardrail's output. Routing asked for no - compression, so it must read the pre-compression snapshot instead.""" - original = [{"role": "user", "content": "the full original conversation"}] - auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original)) + async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self): + """Routing asked for no compression while the model hop compressed, so the + messages in hand are that guardrail's output and no uncompressed copy survives. + Routing reads them as-is: the alternative is keeping a pre-guardrail copy, which + is the text a masking guardrail exists to remove.""" policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] - result = await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) - - assert result == original + assert await messages_for_routing(policy=policy, messages=model_compressed, request_kwargs={}) is None @pytest.mark.asyncio async def test_unknown_guardrail_name_routes_on_the_uncompressed_messages(self): @@ -344,15 +334,18 @@ class TestMessagesForRouting: assert result == [{"role": "user", "content": "[COMPRESSED] hello world"}] @pytest.mark.asyncio - async def test_uses_the_snapshot_when_present(self, registered_guardrail): + async def test_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail): + """Regression (security): routing-side compression POSTs its input to an external + service, so it must read the live messages every earlier guardrail has already + rewritten. Routing on a pre-guardrail copy would send a masking guardrail's own + input straight back out of the proxy.""" policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") - auto_router_compression._routing_messages_snapshot.set(({"role": "user", "content": "original"},)) - # `messages` here stands in for whatever the model-side guardrail already - # rewrote `data["messages"]` to -- routing must ignore it and compress the - # pristine snapshot instead. - already_rewritten = [{"role": "user", "content": "rewritten by another guardrail"}] - result = await messages_for_routing(policy=policy, messages=already_rewritten, request_kwargs={}) - assert result == [{"role": "user", "content": "[COMPRESSED] original"}] + masked = [{"role": "user", "content": "my ssn is [REDACTED]"}] + + result = await messages_for_routing(policy=policy, messages=masked, request_kwargs={}) + + assert result == [{"role": "user", "content": "[COMPRESSED] my ssn is [REDACTED]"}] + assert registered_guardrail.request_data_seen[0]["messages"] == masked @pytest.mark.asyncio async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4c5813911ac..b4de20e9f0a 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10106,32 +10106,29 @@ class TestAutoRouterCompressionDecoupling: assert registered_guardrail.call_count == 0 @pytest.mark.asyncio - async def test_routing_none_routes_on_the_original_not_the_model_compressed_messages( + async def test_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy( self, registered_guardrail ): - """Regression: with routing explicitly off and the model side compressed, the - messages the router holds are the model-side guardrail's output. Routing asked - for no compression, so it has to classify on the pre-compression snapshot.""" - from litellm.proxy.guardrails import auto_router_compression + """Routing asked for no compression while the model hop compressed, so the only + messages left are that guardrail's output and the strategy classifies on them. + Keeping a pre-compression copy to classify on instead is what this deliberately + gives up: that copy is taken before the pre-call guardrails run, so it still + holds whatever a masking guardrail exists to strip, and routing-side compression + POSTs its input to an external service.""" router, strategy = self._router( { "auto_router_routing_compression": "none", "auto_router_model_compression": "fake-compress", } ) - original_messages = self._messages() - auto_router_compression._routing_messages_snapshot.set(tuple(dict(m) for m in original_messages)) model_compressed = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] - try: - await router.async_pre_routing_hook( - model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed - ) - finally: - auto_router_compression._routing_messages_snapshot.set(None) + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=model_compressed + ) - assert strategy.received_messages == original_messages + assert strategy.received_messages == model_compressed assert registered_guardrail.call_count == 0 @pytest.mark.asyncio From 14f8677bfcdc160d3b3b424dc84a9c1727734939 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:57:51 -0700 Subject: [PATCH 22/61] fix(realtime): mark realtime sessions async so failure hooks fire once The relay's failure dispatch runs the async handler and then the legacy sync failure_handler for the proxy's callable callbacks. The realtime logging object carried no async marker, so failure_handler treated the session as a sync SDK call and fired every CustomLogger's sync failure hook on top of the async one: Langfuse recorded two ERROR observations per refused session, and OpenTelemetry, MLflow, Braintrust, Literal AI, DeepEval and New Relic implement the same sync hook. Plant the _arealtime marker in litellm_params the way aanthropic_messages and agenerate_content already do, so both dispatchers classify the session async. --- litellm/litellm_core_utils/litellm_logging.py | 1 + litellm/realtime_api/main.py | 4 +-- .../test_litellm_logging.py | 30 +++++++++++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index df579f6df5b..15585c64efb 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1820,6 +1820,7 @@ class Logging(LiteLLMLoggingBaseClass): and litellm_params.get(CallTypes.aanthropic_messages.value, False) is not True and litellm_params.get(CallTypes.agenerate_content.value, False) is not True and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True + and litellm_params.get(CallTypes.arealtime.value, False) is not True ) def _is_assembled_stream_success(self, result=None) -> bool: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 3862aec445f..9de91dfcaa5 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -27,7 +27,7 @@ from litellm.types.realtime import ( RealtimeTranscriptionSessionRequest, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LlmProviders +from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager from ..litellm_core_utils.get_litellm_params import get_litellm_params @@ -355,7 +355,7 @@ async def _arealtime( user: Final = kwargs.get("user", None) litellm_params: Final = GenericLiteLLMParams(**kwargs) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.arealtime.value: True} model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider( model=model, diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f1de7390b5b..af75691eb10 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -995,6 +995,35 @@ async def test_anthropic_messages_marks_litellm_params_async(): litellm.callbacks = original_callbacks +@pytest.mark.asyncio +async def test_arealtime_marks_litellm_params_async(monkeypatch): + """LIT-6973: ``_arealtime`` must plant ``_arealtime`` in ``litellm_params`` so + ``_is_sync_litellm_request`` classifies the session async and a failed session + reaches a CustomLogger's failure hook once, through the async path only, even + though the sync ``failure_handler`` still runs ahead of the async one.""" + captured = {} + async_logged = asyncio.Event() + + class CaptureLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + captured["litellm_params"] = kwargs.get("litellm_params", {}) + async_logged.set() + + logger = CaptureLogger() + logger.log_failure_event = MagicMock() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + with pytest.raises(ValueError, match="Unsupported model"): + await litellm._arealtime(model="anthropic/claude-x", websocket=MagicMock()) + await asyncio.wait_for(async_logged.wait(), timeout=10) + logger.log_failure_event.assert_not_called() + assert captured["litellm_params"].get("_arealtime") is True + assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False + + @pytest.mark.asyncio async def test_agenerate_content_marks_litellm_params_async(): """LIT-4475: the async ``agenerate_content`` entrypoint must plant @@ -1180,6 +1209,7 @@ def test_is_sync_litellm_request(): assert LitellmLogging._is_sync_litellm_request({}) is True assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False + assert LitellmLogging._is_sync_litellm_request({"_arealtime": True}) is False assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False From 412c36bb8e0663fd27e6c635d94f35d7407eabd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:24:36 -0700 Subject: [PATCH 23/61] fix(realtime): detect an upstream refusal from received frames, not the session log The refusal predicate also required the session log to be empty, but that log is not limited to upstream frames. With gemini_live_defer_setup the handler stores a synthetic session.created before the relay starts, and the transcription usage flush appends a usage event before the check runs, so an upstream policy close with no received frames was still logged as a $0 success. Key the check off the received-frames flag only --- .../litellm_core_utils/realtime_streaming.py | 2 +- .../test_realtime_streaming.py | 39 +++++++++++++++++-- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index bb7fbd81146..984934daaac 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1135,7 +1135,7 @@ class RealTimeStreaming: return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") def _backend_refused_session(self, close: BackendClose) -> bool: - return close.code != 1000 and not self._backend_sent_frames and not self.messages + return close.code != 1000 and not self._backend_sent_frames async def log_backend_refusal(self, error: Exception) -> None: if not self.logging_obj: diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 00addb613c2..09757c35570 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3031,12 +3031,12 @@ async def test_session_close_flushes_unbilled_transcription_usage(): messages before log_messages runs, and never forwarded to the client.""" from typing import Final - from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage + from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeResponseTypedDict client_ws: Final = MagicMock() client_ws.send_text = AsyncMock() backend_ws: Final = MagicMock() - backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + backend_ws.recv = AsyncMock(side_effect=[b'{"serverContent": {}}', ConnectionClosed(None, None)]) logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() @@ -3048,7 +3048,24 @@ async def test_session_close_flushes_unbilled_transcription_usage(): "total_tokens": 171, "input_token_details": {"text_tokens": 0, "audio_tokens": 153}, } + transcript_frame: Final[RealtimeResponseTypedDict] = { + "response": { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "transcript": "ahoy", + "item_id": "item_1", + "content_index": 0, + }, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } provider_config: Final = MagicMock() + provider_config.transform_realtime_response = MagicMock(return_value=transcript_frame) provider_config.unbilled_usage_on_session_close = MagicMock(return_value=usage) streaming: Final = RealTimeStreaming( @@ -3080,7 +3097,9 @@ async def test_session_close_flushes_unbilled_transcription_usage(): ) assert len(flushed) == 1 assert flushed[0] in logged_snapshots[0] - assert not client_ws.send_text.called + forwarded: Final = tuple(json.loads(call.args[0]) for call in client_ws.send_text.await_args_list) + assert [event.get("transcript") for event in forwarded] == ["ahoy"] + assert all("usage" not in event for event in forwarded) @pytest.mark.asyncio @@ -3246,6 +3265,20 @@ async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success(): assert session.logging.logged_sessions == () +@pytest.mark.asyncio +async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_failure(): + """LIT-6973: deferred Gemini Live setup stores a synthetic ``session.created`` before + the relay starts. It is not an upstream frame, so a refusal after it is still a refusal.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + session.streaming.store_message(json.dumps({"type": "session.created", "session": {"id": "sess_synthetic"}})) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert session.logging.logged_sessions == () + + @pytest.mark.asyncio async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success(): client_ws: Final = _client_ws_that_never_sends() From 74613f9bd47d8e3068e6e2f1f519675ac15b7ab8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:52:26 -0700 Subject: [PATCH 24/61] fix(realtime): redact credentials from the relayed upstream close The handshake error path already runs client-facing error strings through _redact_string; the relay's _close_client did not, so a secret echoed in an upstream close reason could reach the client verbatim. Mirror the handshake path and scrub the close message and reason before relaying them. --- .../litellm_core_utils/realtime_streaming.py | 8 +++++--- .../test_realtime_streaming.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 984934daaac..b448cb7c9ff 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cas from typing_extensions import ReadOnly import litellm -from litellm._logging import verbose_logger +from litellm._logging import _redact_string, verbose_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -1567,12 +1567,14 @@ class RealTimeStreaming: await asyncio.gather(forward_task, client_task, return_exceptions=True) async def _close_client(self, close: BackendClose) -> None: + redacted_message: Final = _redact_string(close.message) + redacted_reason: Final = _redact_string(close.reason) try: if close.code != 1000: - await self.websocket.send_text(realtime_error_event(close.message, error_type="server_error")) + await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error")) await self.websocket.close( code=client_close_code(close.code), - reason=websocket_close_reason(close.reason, fallback=close.message), + reason=websocket_close_reason(redacted_reason, fallback=redacted_message), ) except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way verbose_logger.debug("Could not relay the upstream close to the client: %s", e) diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 09757c35570..bcfacf16205 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3229,6 +3229,24 @@ async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +@pytest.mark.asyncio +async def test_upstream_close_reason_with_a_secret_is_redacted_before_reaching_the_client(): + """LIT-6973: the relayed close mirrors the handshake path and scrubs credential + patterns, so an upstream error echoing a token never reaches the client verbatim.""" + secret: Final = "sk-live-abcdef0123456789abcdef0123" + client_ws: Final = _client_ws_that_never_sends() + upstream_close: Final = ConnectionClosed(Close(1008, f"auth failed for {secret}"), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) + + await session.run() + + (error_event,) = _error_events_sent_to(client_ws) + assert secret not in error_event["error"]["message"] + relayed_reason: Final = client_ws.close.await_args.kwargs["reason"] + assert secret not in relayed_reason + assert "REDACTED" in relayed_reason + + @pytest.mark.asyncio async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error(): client_ws: Final = _client_ws_that_never_sends() From af3ddb477a852f20898aeefd7bc35713188f98da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:22:20 -0700 Subject: [PATCH 25/61] fix(realtime): release the budget reservation on a failed session and scrub relayed close details A refused or failed /v1/realtime session never ran the success cost callback or a failure hook, so its pre-call budget reservation stayed open and kept the key/team/user spend counters pinned above real spend, 429ing later requests on the same key until the counter's TTL expired. The endpoint now reconciles the reservation in a finally, reusing a shared release_or_invalidate_budget_reservation helper that mirrors the success/failure paths (release to zero, else invalidate the reserved counters and finalize). The relayed upstream close message and reason also go through the proxy's client-facing redaction, so a credential, internal hostname, private IP, or server path echoed by the upstream never reaches the client verbatim. --- .../litellm_core_utils/realtime_streaming.py | 6 +- litellm/proxy/proxy_server.py | 12 +++ .../spend_tracking/budget_reservation.py | 25 ++++++ .../test_realtime_streaming.py | 21 +++-- tests/test_litellm/proxy/test_proxy_server.py | 83 +++++++++++++++++++ 5 files changed, 137 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b448cb7c9ff..c670278d3fb 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cas from typing_extensions import ReadOnly import litellm -from litellm._logging import _redact_string, verbose_logger +from litellm._logging import redact_internal_details_from_client_message, verbose_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.types.llms.openai import ( @@ -1567,8 +1567,8 @@ class RealTimeStreaming: await asyncio.gather(forward_task, client_task, return_exceptions=True) async def _close_client(self, close: BackendClose) -> None: - redacted_message: Final = _redact_string(close.message) - redacted_reason: Final = _redact_string(close.reason) + redacted_message: Final = redact_internal_details_from_client_message(close.message) + redacted_reason: Final = redact_internal_details_from_client_message(close.reason) try: if close.code != 1000: await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error")) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 61bcdea94d4..06cfee45918 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11453,6 +11453,16 @@ def _realtime_query_params_template(model: str | None, intent: str | None) -> tu return tuple(params) +async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth) -> None: + from litellm.proxy.spend_tracking.budget_reservation import ( + release_or_invalidate_budget_reservation, + ) + + await release_or_invalidate_budget_reservation( + budget_reservation=user_api_key_dict.budget_reservation, + ) + + @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11592,6 +11602,8 @@ async def realtime_websocket_endpoint( ) except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") + finally: + await _release_realtime_budget_reservation(user_api_key_dict) ###################################################################### diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 91d2ece7a51..2ee7320b82c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -373,6 +373,31 @@ async def invalidate_budget_reservation_counters( await _invalidate_spend_counter(counter_key=counter_key) +async def release_or_invalidate_budget_reservation( + budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict +) -> None: + """Reconcile a still-open reservation on a terminal path that settles no cost. + + A failed or upstream-refused request never runs the success cost callback, so + its pre-call reservation stays open and keeps the spend counter pinned above + real spend until the counter's TTL expires, 429ing later requests on the same + key. Release it to zero; if the release itself fails (e.g. the counter store is + unreachable) drop the reserved counters directly and mark the reservation + finalized so nothing reprocesses it. Idempotent: the finalized guard makes a + second call a no-op once success or failure handling already reconciled. + """ + if budget_reservation is None or budget_reservation.get("finalized") is True: + return + try: + await release_budget_reservation(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead + verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters") + try: + await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + finally: + budget_reservation["finalized"] = True + + async def _get_budget_counters( request_body: dict, valid_token: UserAPIKeyAuth, diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index bcfacf16205..5352c894f87 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3229,21 +3229,28 @@ async def test_bidirectional_forward_relays_upstream_policy_close_to_client(): client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL) +@pytest.mark.parametrize( + "leaked_detail", + ( + pytest.param("sk-live-abcdef0123456789abcdef0123", id="credential"), + pytest.param("vertex-int.svc.cluster.local", id="internal-hostname"), + pytest.param("/etc/litellm/service-account.json", id="filesystem-path"), + ), +) @pytest.mark.asyncio -async def test_upstream_close_reason_with_a_secret_is_redacted_before_reaching_the_client(): - """LIT-6973: the relayed close mirrors the handshake path and scrubs credential - patterns, so an upstream error echoing a token never reaches the client verbatim.""" - secret: Final = "sk-live-abcdef0123456789abcdef0123" +async def test_upstream_close_details_are_scrubbed_before_reaching_the_client(leaked_detail: str): + """LIT-6973: the relayed close goes through the proxy's client-facing redaction, so an upstream + error echoing a credential, an internal host, or a server path never reaches the client verbatim.""" client_ws: Final = _client_ws_that_never_sends() - upstream_close: Final = ConnectionClosed(Close(1008, f"auth failed for {secret}"), None) + upstream_close: Final = ConnectionClosed(Close(1008, f"upstream rejected: {leaked_detail}"), None) session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close)) await session.run() (error_event,) = _error_events_sent_to(client_ws) - assert secret not in error_event["error"]["message"] + assert leaked_detail not in error_event["error"]["message"] relayed_reason: Final = client_ws.close.await_args.kwargs["reason"] - assert secret not in relayed_reason + assert leaked_detail not in relayed_reason assert "REDACTED" in relayed_reason diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d91928a203e..57e2cfc3332 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9521,6 +9521,89 @@ def test_realtime_websocket_route_aliases_registered(): ) +def _lit6973_fake_realtime_ws() -> MagicMock: + ws = MagicMock() + ws.headers = {} + ws.scope = {"headers": [], "type": "websocket"} + ws.url = "ws://testserver/v1/realtime" + ws.accept = AsyncMock() + ws.send_text = AsyncMock() + ws.close = AsyncMock() + return ws + + +async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: + """Drive realtime_websocket_endpoint through a session the upstream refused. + + route_request resolves normally because the relay handles the refusal + internally (sends the error event, closes the client), so neither the + success cost callback nor a failure hook runs on _ProxyDBLogger. The + endpoint itself must reconcile the pre-call budget reservation, so the + real release runs (entries is empty, so it touches no counter store) and + the caller asserts on the observable reservation state afterwards.""" + from litellm.proxy import proxy_server as ps + + user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") + user_api_key_dict.budget_reservation = reservation + + completed: Final = asyncio.get_running_loop().create_future() + completed.set_result(None) + + pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, MagicMock())) + can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test + pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state + route = patch.object(ps, "route_request", new=AsyncMock(return_value=completed)) # test-quality-ok: fakes the relay that already handled the refusal so the session returns normally + with can_call, pre, route: + await ps.realtime_websocket_endpoint( + websocket=_lit6973_fake_realtime_ws(), + model="vertex_ai/gemini-live-2.5-flash", + intent=None, + guardrails=None, + user_api_key_dict=user_api_key_dict, + ) + + +@pytest.mark.asyncio +async def test_refused_realtime_session_releases_the_budget_reservation(): + """LIT-6973: reclassifying a refused realtime session as a failure removed the + success-path reservation release, so the pre-call reservation stayed open and + pinned the key/team/user spend counters, locking the key after a couple of + refusals. The endpoint must reconcile it: the reservation ends up finalized.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + await _lit6973_drive_refused_realtime_session(reservation) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): + """If releasing the reservation itself fails (e.g. the counter store is down), + the reserved counters must be invalidated directly so the estimate does not + stay pinned, and the reservation is finalized so nothing reprocesses it.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = { + "reserved_cost": 0.55, + "input_cost": 0.0, + "finalized": False, + "entries": [{"counter_key": "spend:key:hashed-token"}], + } + invalidated: Final[list[str]] = [] + + async def _record(counter_key: str) -> None: + invalidated.append(counter_key) + + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated + sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable + with failing_release, sink: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert invalidated == ["spend:key:hashed-token"] + assert reservation["finalized"] is True + + class TestTransformRequestBannedParams: """ /utils/transform_request applies the same banned-param check as LLM endpoints. From 1fe87e8e25206c039be5e87a5808a30f47cc3183 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:11:24 -0700 Subject: [PATCH 26/61] fix(realtime): settle the budget reservation only for sessions the success log does not own The blanket finally release from the previous commit also zeroed the reservation of successful sessions. Success settlement is enqueued on the logging worker, not awaited, so the endpoint's finally ran first and released the reservation the cost callback still had to reconcile, dropping the real spend from the key/team/user counters. The relay now stamps a synchronous marker (REALTIME_SESSION_SUCCESS_LOGGED_KEY) on the shared logging object at the single success-dispatch site, and the endpoint releases the reservation only when that marker is absent. Refused or failed sessions, which never log success, still release; successful sessions leave the reservation for the cost callback to settle to actual spend. Exactly one settler touches each reservation, so the idempotent reconcile never double-adjusts. --- .../litellm_core_utils/realtime_streaming.py | 4 ++ litellm/proxy/proxy_server.py | 7 ++- .../test_realtime_streaming.py | 32 +++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 55 +++++++++++++------ 4 files changed, 80 insertions(+), 18 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c670278d3fb..75046f2cf87 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -35,6 +35,9 @@ else: CLIENT_CONNECTION_CLASS = Any +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" + + @dataclass(frozen=True, slots=True) class BackendClose: code: int @@ -421,6 +424,7 @@ class RealTimeStreaming: self._logging_worker.ensure_initialized_and_enqueue( self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) ) + self.logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True async def _send_to_backend(self, message: str) -> bool: """Send a message to the backend WebSocket. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06cfee45918..7d59dfa86c4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11603,7 +11603,12 @@ async def realtime_websocket_endpoint( except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone") finally: - await _release_realtime_budget_reservation(user_api_key_dict) + from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + ) + + if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): + await _release_realtime_budget_reservation(user_api_key_dict) ###################################################################### diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 5352c894f87..9c0f6f59463 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -14,6 +14,7 @@ import litellm from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, RealTimeStreaming, client_sent_openai_beta_realtime_header, ) @@ -3380,3 +3381,34 @@ async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the assert session.logging.logged_sessions == ((),) assert session.logging.logged_failures == () client_ws.close.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_success_logging_stamps_the_reservation_ownership_marker(): + """LIT-6973: only the success path enqueues the cost callback that settles the + session's budget reservation, so it stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on + the shared logging object. The proxy endpoint reads that stamp to decide whether to + release the reservation itself, so a logged-as-success session must carry it.""" + client_ws: Final = _client_ws_that_never_sends() + session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() + upstream_close: Final = ConnectionClosed(Close(1000, ""), None) + session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close)) + + await session.run() + + assert session.logging.logged_sessions != () + assert session.logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + +@pytest.mark.asyncio +async def test_refused_session_does_not_stamp_the_reservation_ownership_marker(): + """A refused session logs a failure, not a success, so it must not stamp + REALTIME_SESSION_SUCCESS_LOGGED_KEY. If it did, the proxy endpoint would skip its + own reservation release and the refused session's reservation would stay pinned.""" + upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None) + session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close)) + + await session.run() + + assert session.logging.logged_failures == (upstream_close,) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 57e2cfc3332..697d4c182c7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9532,27 +9532,34 @@ def _lit6973_fake_realtime_ws() -> MagicMock: return ws -async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: - """Drive realtime_websocket_endpoint through a session the upstream refused. +async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_success: bool) -> None: + """Drive realtime_websocket_endpoint to just before its budget-reservation finally. - route_request resolves normally because the relay handles the refusal - internally (sends the error event, closes the client), so neither the - success cost callback nor a failure hook runs on _ProxyDBLogger. The - endpoint itself must reconcile the pre-call budget reservation, so the - real release runs (entries is empty, so it touches no counter store) and - the caller asserts on the observable reservation state afterwards.""" + route_request resolves normally in both cases: the relay owns the session + once route_request returns. A successful session enqueues its success cost + callback and stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on the shared logging + object; a refused one does neither. The endpoint keys its reservation cleanup + off that stamp, so backend_logged_success reproduces both branches. The fake + logging object carries a real model_call_details dict so the stamp is + observable, and the reservation has empty entries so the real release touches + no counter store.""" + from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.proxy import proxy_server as ps user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token") user_api_key_dict.budget_reservation = reservation - completed: Final = asyncio.get_running_loop().create_future() - completed.set_result(None) + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} - pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, MagicMock())) + async def fake_llm_call() -> None: + if backend_logged_success: + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + + pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)) can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state - route = patch.object(ps, "route_request", new=AsyncMock(return_value=completed)) # test-quality-ok: fakes the relay that already handled the refusal so the session returns normally + route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( websocket=_lit6973_fake_realtime_ws(), @@ -9565,17 +9572,31 @@ async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: @pytest.mark.asyncio async def test_refused_realtime_session_releases_the_budget_reservation(): - """LIT-6973: reclassifying a refused realtime session as a failure removed the - success-path reservation release, so the pre-call reservation stayed open and - pinned the key/team/user spend counters, locking the key after a couple of - refusals. The endpoint must reconcile it: the reservation ends up finalized.""" + """LIT-6973: a refused realtime session enqueues no success cost callback, so + the pre-call reservation would stay open and pin the key/team/user spend + counters, locking the key after a couple of refusals. The endpoint sees no + success stamp and reconciles it: the reservation ends up finalized.""" reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} - await _lit6973_drive_refused_realtime_session(reservation) + await _lit6973_drive_realtime_session(reservation, backend_logged_success=False) assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): + """A billable realtime session settles its reservation through the enqueued + success cost callback, not the endpoint. The endpoint must not finalize it in + its finally, or it would reconcile the reservation to zero before the cost + callback applies real spend, so billable sessions stop counting against budget. + With the success stamp present, the endpoint leaves the reservation untouched.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + await _lit6973_drive_realtime_session(reservation, backend_logged_success=True) + + assert reservation["finalized"] is False + + @pytest.mark.asyncio async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): """If releasing the reservation itself fails (e.g. the counter store is down), From 5a35e6d41f76d2b09a258b3e2b051e3e7e745c79 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:18:33 -0700 Subject: [PATCH 27/61] fix(realtime): release the budget reservation when a session is rejected before the relay starts The three pre-relay exits of realtime_websocket_endpoint (missing model, key/model access denied, pre-call rejection such as a rate limit or a guardrail) returned before the finally that releases the auth-time budget reservation, so a rejected session pinned the key at the reserved amount until the counter TTL expired and its next requests got budget_exceeded while /key/info showed spend 0. A single _reject_realtime_session helper now releases the reservation before sending the error event and closing, and release_or_invalidate_budget_reservation shields the release from a second cancellation and logs, rather than raises, a failing invalidate fallback so it can never mask the session's own outcome. --- litellm/proxy/proxy_server.py | 43 ++++++---- .../spend_tracking/budget_reservation.py | 4 +- tests/test_litellm/proxy/test_proxy_server.py | 79 +++++++++++++++++-- 3 files changed, 103 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d59dfa86c4..65fe3ede822 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11463,6 +11463,25 @@ async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth ) +async def _reject_realtime_session( + websocket: WebSocket, + user_api_key_dict: UserAPIKeyAuth, + *, + code: int, + reason: str, + error_message: str | None = None, +) -> None: + await _release_realtime_budget_reservation(user_api_key_dict) + if error_message is not None: + try: + await websocket.send_text( + json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) + ) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") + await websocket.close(code=code, reason=reason) + + @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11488,7 +11507,9 @@ async def realtime_websocket_endpoint( if intent == "transcription": route_model = "gpt-realtime-whisper" else: - await websocket.close(code=1008, reason="model query parameter is required") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1008, reason="model query parameter is required" + ) return assert route_model is not None try: @@ -11499,7 +11520,7 @@ async def realtime_websocket_endpoint( llm_router=llm_router, ) except ProxyException as e: - await websocket.close(code=1008, reason=e.message[:120]) + await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120]) return await websocket.accept(**accept_kwargs) @@ -11558,21 +11579,9 @@ async def realtime_websocket_endpoint( ) except Exception as e: verbose_proxy_logger.exception("Realtime pre-call error") - try: - await websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "guardrail_error", - "message": str(e), - }, - } - ) - ) - except Exception: - pass - await websocket.close(code=1011, reason="Pre-call error") + await _reject_realtime_session( + websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e) + ) return # Phase 2: route to upstream LLM. diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 2ee7320b82c..ed2bc87597c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -389,11 +389,13 @@ async def release_or_invalidate_budget_reservation( if budget_reservation is None or budget_reservation.get("finalized") is True: return try: - await release_budget_reservation(budget_reservation=budget_reservation) + await asyncio.shield(release_budget_reservation(budget_reservation=budget_reservation)) except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters") try: await invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed + verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after release failed") finally: budget_reservation["finalized"] = True diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 697d4c182c7..8b151d9e1f6 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9532,8 +9532,15 @@ def _lit6973_fake_realtime_ws() -> MagicMock: return ws -async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_success: bool) -> None: - """Drive realtime_websocket_endpoint to just before its budget-reservation finally. +async def _lit6973_drive_realtime_session( + reservation: dict, *, backend_logged_success: bool, phase_one_exit: str | None = None +) -> MagicMock: + """Drive realtime_websocket_endpoint through one of its reservation-settling exits. + + phase_one_exit picks a rejection before the relay: "model_access" makes the + key/model check raise ProxyException, "pre_call" makes pre-call processing + (rate limits, guardrails) raise. Neither reaches route_request, so no success + log can own the reservation and the endpoint has to release it on that exit. route_request resolves normally in both cases: the relay owns the session once route_request returns. A successful session enqueues its success cost @@ -9556,18 +9563,30 @@ async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_s if backend_logged_success: logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True - pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)) - can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally under test + from litellm.proxy._types import ProxyException + + model_access_error: Final = ( + ProxyException(message="key cannot access model", type="auth_error", param="model", code=401) + if phase_one_exit == "model_access" + else None + ) + pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None + pre_call: Final = AsyncMock( + side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) + ) + ws: Final = _lit6973_fake_realtime_ws() + can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object with can_call, pre, route: await ps.realtime_websocket_endpoint( - websocket=_lit6973_fake_realtime_ws(), + websocket=ws, model="vertex_ai/gemini-live-2.5-flash", intent=None, guardrails=None, user_api_key_dict=user_api_key_dict, ) + return ws @pytest.mark.asyncio @@ -9583,6 +9602,39 @@ async def test_refused_realtime_session_releases_the_budget_reservation(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reservation(): + """A rate-limit or guardrail rejection happens before route_request, so the + relay never runs and no success log can own the reservation. The endpoint + must release it on that exit too, or the key stays pinned at the reserved + amount and its next requests 429 with budget_exceeded while /key/info shows + spend 0 (reproduced live with rpm_limit=1). The client still gets the + pre-call error event and the 1011 close it got before.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call" + ) + + assert reservation["finalized"] is True + assert json.loads(ws.send_text.await_args.args[0])["error"]["message"] == "Rate limit exceeded" + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + + +@pytest.mark.asyncio +async def test_realtime_session_denied_model_access_releases_the_budget_reservation(): + """The key/model access check rejects before the socket is even accepted; + that exit skipped the release as well, pinning the reservation.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + + ws: Final = await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access" + ) + + assert reservation["finalized"] is True + ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") + + @pytest.mark.asyncio async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): """A billable realtime session settles its reservation through the enqueued @@ -9625,6 +9677,23 @@ async def test_release_or_invalidate_falls_back_to_invalidating_the_counters(): assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback_fails(): + """Both counter-store calls failing must not raise out of the realtime + endpoint's finally (it would mask the session's own outcome) and must still + stamp finalized so nothing retries the same reservation.""" + from litellm.proxy.spend_tracking import budget_reservation as br + + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch + failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail + + with failing_release, failing_invalidate: + await br.release_or_invalidate_budget_reservation(budget_reservation=reservation) + + assert reservation["finalized"] is True + + class TestTransformRequestBannedParams: """ /utils/transform_request applies the same banned-param check as LLM endpoints. From 37722eba68149c5f3e59ed0f3c12798a84aa1bc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:35:26 -0700 Subject: [PATCH 28/61] fix(realtime): close a rejected client before releasing its budget reservation A slow or unreachable counter store made a pre-relay rejection wait behind the reservation release before the client saw the error event and the close. Close first and release in finally, mirroring the relay's own failure path, so a client that already hung up still gets its reservation released. --- litellm/proxy/proxy_server.py | 20 ++++---- tests/test_litellm/proxy/test_proxy_server.py | 47 ++++++++++++++++++- 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 65fe3ede822..a5c60d8e976 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11471,15 +11471,17 @@ async def _reject_realtime_session( reason: str, error_message: str | None = None, ) -> None: - await _release_realtime_budget_reservation(user_api_key_dict) - if error_message is not None: - try: - await websocket.send_text( - json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) - ) - except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below - verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") - await websocket.close(code=code, reason=reason) + try: + if error_message is not None: + try: + await websocket.send_text( + json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}}) + ) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway") + await websocket.close(code=code, reason=reason) + finally: + await _release_realtime_budget_reservation(user_api_key_dict) @app.websocket("/openai/v1/realtime") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8b151d9e1f6..4d70c9d436f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9533,7 +9533,11 @@ def _lit6973_fake_realtime_ws() -> MagicMock: async def _lit6973_drive_realtime_session( - reservation: dict, *, backend_logged_success: bool, phase_one_exit: str | None = None + reservation: dict, + *, + backend_logged_success: bool, + phase_one_exit: str | None = None, + websocket: MagicMock | None = None, ) -> MagicMock: """Drive realtime_websocket_endpoint through one of its reservation-settling exits. @@ -9574,7 +9578,7 @@ async def _lit6973_drive_realtime_session( pre_call: Final = AsyncMock( side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj) ) - ws: Final = _lit6973_fake_realtime_ws() + ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws() can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object @@ -9635,6 +9639,45 @@ async def test_realtime_session_denied_model_access_releases_the_budget_reservat ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model") +@pytest.mark.asyncio +async def test_rejected_realtime_session_closes_the_client_before_releasing_the_reservation(): + """The counter release can block on a slow or unreachable store, and a + rejected client must not sit behind it: the relay's own failure path closes + the client first and releases in its finally, so the pre-relay rejection + has to close first as well. The fake close checks the reservation is still + open when the client is closed.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + + async def close_while_reservation_is_still_open(**_: object) -> None: + assert reservation["finalized"] is False, "client was closed only after the reservation release" + + ws.close = AsyncMock(side_effect=close_while_reservation_is_still_open) + + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="pre_call", websocket=ws + ) + + ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error") + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_rejected_realtime_session_releases_the_reservation_when_the_client_is_already_gone(): + """A client that hung up before the rejection makes the close raise; the + reservation must still be released, or the key stays pinned.""" + reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + ws: Final = _lit6973_fake_realtime_ws() + ws.close = AsyncMock(side_effect=RuntimeError("client already disconnected")) + + with pytest.raises(RuntimeError, match="client already disconnected"): + await _lit6973_drive_realtime_session( + reservation, backend_logged_success=False, phase_one_exit="model_access", websocket=ws + ) + + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback(): """A billable realtime session settles its reservation through the enqueued From 03da725ee4de2414056765f1968794e4c0634ce2 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:30:42 -0700 Subject: [PATCH 29/61] Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/test_litellm/integrations/test_shadow_eval_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index d67e957ca16..371f6f75a05 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -159,7 +159,7 @@ def _reasoning_judge_router( if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN: kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"} return {"choices": [{"message": {"content": "shadow answer"}}]} - budget_for_the_answer = kwargs["max_tokens"] - reasoning_tokens + budget_for_the_answer: Final = kwargs["max_tokens"] - reasoning_tokens return {"choices": [{"message": {"content": verdict[: max(0, budget_for_the_answer)]}}]} router.acompletion = MagicMock(side_effect=acompletion) From 00b49ccc8bb0de73891376e5ee1e8bd58295ba2d Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:31:18 -0700 Subject: [PATCH 30/61] fix(auto-router compression): honor the policy on the SDK path and on re-save The router reused the model hop's compression for routing whenever both hops named the same guardrail, on the premise that arm_pre_call had already run it. Only the proxy calls arm_pre_call, so through the SDK nothing armed the guardrail and nothing had compressed anything: the shortcut skipped routing compression too and served the request with no compression on either hop. The reuse is now conditional on the model hop actually having been armed. The Admin UI hydrated an absent auto_router_model_compression as same-as-routing, while the backend reads it as no model-hop compression. Opening a router configured with only auto_router_routing_compression and saving any unrelated edit wrote the routing guardrail onto the model hop, silently starting to compress the model call. Both carry a regression test that fails when the fix is reverted. --- .../guardrails/auto_router_compression.py | 15 + litellm/router.py | 8 +- tests/test_litellm/test_router.py | 967 ++++++------------ .../buildAutoRouterCompression.test.ts | 15 +- .../add_model/buildAutoRouterCompression.ts | 8 +- 5 files changed, 347 insertions(+), 666 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 2ab779658b5..e7f58662249 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -45,6 +45,19 @@ def suppressed_compression_guardrails() -> frozenset[str]: return _suppressed_compression_guardrails.get() +# Whether `arm_pre_call` actually armed a model-side compression guardrail for this +# request. Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and +# nothing compresses; the router must not assume the model hop already ran. +_model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( + "litellm_auto_router_model_hop_armed", default=False +) + + +def model_hop_compression_armed() -> bool: + """True when this request's model-side compression guardrail was actually armed.""" + return _model_hop_armed.get() + + @dataclass(frozen=True, slots=True) class AutoRouterCompressionPolicy: """An auto router's compression choice for each hop. ``None`` means no compression.""" @@ -147,6 +160,7 @@ async def arm_pre_call( guardrail the policy names (if any) even when it isn't ``default_on``. """ _suppressed_compression_guardrails.set(frozenset()) + _model_hop_armed.set(False) if llm_router is None: return @@ -179,6 +193,7 @@ async def arm_pre_call( ) if policy.model is not None: + _model_hop_armed.set(True) _, metadata = get_or_create_metadata_bucket(data) requested: Final = metadata.get("guardrails") existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () diff --git a/litellm/router.py b/litellm/router.py index 989914b1610..b61e4e29e09 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13039,6 +13039,7 @@ class Router: from litellm.proxy.guardrails.auto_router_compression import ( messages_for_routing, + model_hop_compression_armed, policy_for_model, team_id_from_request, ) @@ -13057,8 +13058,13 @@ class Router: # (arm_pre_call armed it whether or not it is `default_on`); reuse that result # for routing too instead of paying for a second compression call against the # same content. + # + # Only the proxy calls arm_pre_call, so that reuse is conditional on it having + # actually run: on the SDK path nothing arms the model hop and nothing has + # compressed anything, and taking the shortcut there would skip both hops and + # silently serve the request with no compression at all. needs_independent_routing_compression: Final = compression_policy is not None and not ( - compression_policy.is_same and compression_policy.model is not None + compression_policy.is_same and compression_policy.model is not None and model_hop_compression_armed() ) routing_messages: Final = ( await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a59c9c98163..9e6a88d3433 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15,7 +15,6 @@ import pytest import respx - import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError @@ -137,31 +136,18 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert ( - encrypted_content_callbacks[0].model_group_affinity_config - == model_group_affinity_config - ) - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.callbacks.index(deployment_callback) - ) + assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [ - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ] + encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -192,13 +178,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -239,10 +221,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -262,9 +241,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -311,18 +288,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -333,9 +302,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -939,9 +906,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -962,9 +927,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -1025,7 +988,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1109,9 +1072,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1119,9 +1080,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1137,30 +1098,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1169,9 +1118,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1221,9 +1168,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1297,7 +1242,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1349,15 +1294,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1392,7 +1331,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match='No deployment available') as exc_info: + with pytest.raises(Exception, match="No deployment available") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1420,15 +1359,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1457,16 +1390,10 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with pytest.raises(Exception, match='Mock failure') as exc_info: + with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: + with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with pytest.raises(Exception, match="Mock failure") as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -1534,9 +1461,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( + f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" + ) @pytest.mark.asyncio @@ -1642,14 +1569,10 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] @@ -1744,9 +1667,7 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -1774,9 +1695,7 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -1846,9 +1765,7 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object( - router, "get_deployment_model_info", side_effect=_model_info_with_str_costs - ): + with patch.object(router, "get_deployment_model_info", side_effect=_model_info_with_str_costs): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -1894,9 +1811,7 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object( - router, "get_deployment_model_info", side_effect=Exception("not found") - ): + with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -2156,6 +2071,7 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] + async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -2849,11 +2765,7 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) + base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator class _Iter(base): def __init__(self): @@ -2933,9 +2845,7 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) + router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -3018,9 +2928,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" + assert "model_group" not in fbk.get("metadata", {}), ( + "model_group leaked into 'metadata' instead of 'litellm_metadata'" + ) @pytest.mark.asyncio @@ -3138,9 +3048,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) + fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -3149,9 +3057,7 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), + return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), ), patch.object( router, @@ -3452,9 +3358,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3482,14 +3386,10 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3516,14 +3416,10 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3550,9 +3446,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) counting_threads = [] monkeypatch.setattr( @@ -3648,14 +3542,10 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3683,9 +3573,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counting_threads = [] @@ -3720,9 +3608,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3758,9 +3644,7 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) + monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) count = await router._acount_pre_call_check_tokens( model="m", @@ -3788,9 +3672,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3815,9 +3697,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3859,9 +3739,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -3930,9 +3808,7 @@ def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwarg prompt_only_tokens = router._count_pre_call_check_tokens( messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -4052,7 +3928,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -4067,9 +3943,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -4159,9 +4033,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -4169,15 +4041,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -4188,26 +4056,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -4226,9 +4086,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -4248,9 +4106,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -4283,9 +4139,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -4294,12 +4148,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -4322,9 +4172,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -4360,15 +4208,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -4415,18 +4259,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -4436,29 +4276,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -4495,10 +4323,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( + f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" + ) # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -4510,10 +4337,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( + f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" + ) # Test Case 3: Bedrock converse endpoint kwargs = { @@ -4525,9 +4351,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" + ) # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -4538,9 +4364,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( + f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" + ) def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -4592,14 +4418,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -4691,15 +4513,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -4726,9 +4543,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -4768,9 +4583,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -4806,9 +4619,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -4851,9 +4662,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -4888,15 +4697,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -4937,15 +4742,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -4972,16 +4773,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_deployment_usable_by_team_helpers(): @@ -5021,9 +4814,7 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) + picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -5033,12 +4824,7 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) + assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -5070,9 +4856,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -5084,21 +4868,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -5171,22 +4945,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -5375,9 +5140,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai-gpt", @@ -5456,12 +5219,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -5486,9 +5245,7 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -5517,9 +5274,7 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -5546,9 +5301,7 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -5586,9 +5339,7 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -5619,9 +5370,7 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -5650,9 +5399,7 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -5754,12 +5501,8 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -5791,12 +5534,8 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -5909,6 +5648,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -5935,6 +5675,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] + async def _drain(): async for chunk in result: collected.append(chunk) @@ -6151,23 +5892,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -6212,12 +5947,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -6342,9 +6073,7 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] + deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -6389,9 +6118,7 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6466,9 +6193,7 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6535,9 +6260,7 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): + def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6592,9 +6315,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) + router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) is None ) assert ( @@ -6616,10 +6337,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) + pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -6645,10 +6364,8 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) + default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None ) assert default_result is not None resolved_model, default_deployment = default_result @@ -6656,10 +6373,7 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) + assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" def _router_with_two_deployments(blocked_flags): @@ -6707,10 +6421,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -6781,9 +6492,7 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -6792,9 +6501,7 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) + healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -6811,9 +6518,7 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -6855,9 +6560,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -6866,9 +6569,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -6892,9 +6593,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -6905,12 +6604,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -6944,16 +6638,9 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() + assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=missing_blocked) - ) - is False - ) - assert ( - litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) - ) + litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) is True ) @@ -6993,9 +6680,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -7013,22 +6698,16 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): + def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): + def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -7044,22 +6723,13 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): + def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) + assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 # --------------------------------------------------------------------------- @@ -7368,9 +7038,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -7379,12 +7047,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -7399,12 +7062,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -7461,13 +7119,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -7511,12 +7169,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -7589,13 +7243,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -7654,13 +7303,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_display_name_treats_malformed_values_as_absent(): @@ -7893,13 +7537,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -7928,13 +7575,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7968,7 +7615,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -7984,13 +7633,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7999,9 +7648,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -8015,7 +7662,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -8050,9 +7699,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -8357,9 +8004,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -8468,9 +8113,7 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) + actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) assert actual is expected, params["model"] @@ -8657,22 +8300,16 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) + router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") assert len(router.model_list) == 1 @@ -9438,18 +9075,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -9572,9 +9205,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -9613,9 +9244,9 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] assert breadcrumbs, "no retry breadcrumbs were recorded" - assert any( - "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs - ), "no breadcrumb carried router walk state, so this test cannot see the leak" + assert any("fallback_depth" in breadcrumb for breadcrumb in breadcrumbs), ( + "no breadcrumb carried router walk state, so this test cannot see the leak" + ) for breadcrumb in breadcrumbs: assert "attempted_targets" not in breadcrumb @@ -9661,7 +9292,9 @@ async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_ke breadcrumbs = metadata["previous_models"] assert breadcrumbs, "no retry breadcrumbs were recorded" dumped = json.dumps(breadcrumbs, default=str) - assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + assert container_key in dumped, ( + "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + ) assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -9766,9 +9399,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -9834,9 +9465,7 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] + affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -9978,9 +9607,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -10024,9 +9651,7 @@ class TestAutoRouterCompressionDecoupling: async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): self.call_count += 1 structured_messages = inputs.get("structured_messages") or [] - compressed = [ - {**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages - ] + compressed = [{**m, "content": f"[COMPRESSED] {m.get('content')}"} for m in structured_messages] return {**inputs, "structured_messages": compressed} @staticmethod @@ -10105,9 +9730,7 @@ class TestAutoRouterCompressionDecoupling: assert registered_guardrail.call_count == 0 @pytest.mark.asyncio - async def test_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy( - self, registered_guardrail - ): + async def test_routing_none_classifies_on_the_live_messages_not_a_pre_guardrail_copy(self, registered_guardrail): """Routing asked for no compression while the model hop compressed, so the only messages left are that guardrail's output and the strategy classifies on them. @@ -10131,11 +9754,37 @@ class TestAutoRouterCompressionDecoupling: assert registered_guardrail.call_count == 0 @pytest.mark.asyncio + @pytest.mark.asyncio + async def test_same_compression_still_compresses_routing_when_nothing_armed_it(self, registered_guardrail): + """Regression: only the proxy calls arm_pre_call. Used through the SDK, nothing + arms the model-side guardrail and nothing has compressed anything, so reusing a + model-hop result that was never produced would serve the request with no + compression on either hop, silently ignoring the configuration.""" + from litellm.proxy.guardrails import auto_router_compression + + router, strategy = self._router( + { + "auto_router_routing_compression": "fake-compress", + "auto_router_model_compression": "fake-compress", + } + ) + uncompressed = self._messages() + assert auto_router_compression.model_hop_compression_armed() is False + + await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=uncompressed + ) + + assert strategy.received_messages != uncompressed + assert registered_guardrail.call_count == 1 + async def test_same_compression_on_both_hops_compresses_once(self, registered_guardrail): """The same/different distinction exists so a shared choice does not pay for compression twice: by the time the router runs, `messages` already reflects whatever the ordinary pre-call guardrail pipeline did for the model call, so the routing decision must reuse it rather than calling the guardrail again.""" + from litellm.proxy.guardrails import auto_router_compression + router, strategy = self._router( { "auto_router_routing_compression": "fake-compress", @@ -10145,13 +9794,17 @@ class TestAutoRouterCompressionDecoupling: # Stands in for what the proxy's ordinary pre-call guardrail pipeline would # have already produced for the model call, since `auto_router_model_compression` # names a guardrail: the router never triggers that pipeline itself. - already_compressed_messages = [ - {"role": "user", "content": "[COMPRESSED] What is the capital of France?"} - ] + already_compressed_messages = [{"role": "user", "content": "[COMPRESSED] What is the capital of France?"}] + # arm_pre_call is what would have armed that guardrail, and only the proxy calls + # it; the reuse below is conditional on it having run. + armed = auto_router_compression._model_hop_armed.set(True) - response = await router.async_pre_routing_hook( - model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages - ) + try: + response = await router.async_pre_routing_hook( + model="smart-router", request_kwargs={"metadata": {}}, messages=already_compressed_messages + ) + finally: + auto_router_compression._model_hop_armed.reset(armed) assert strategy.received_messages == already_compressed_messages assert response.messages == already_compressed_messages @@ -10197,17 +9850,14 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" + assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + f"unexpected error log: {mock_error.call_args_list}" + ) # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -10215,17 +9865,14 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: + with patch("litellm.router.verbose_router_logger.error") as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" + assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( + "expected the error log for an unmappable azure deployment name" + ) # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -10252,6 +9899,7 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -10342,7 +9990,6 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o assert result.supported_reasoning_efforts is None - def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose registry entry declares parallel function calling must flip the group to True instead of False.""" @@ -10575,6 +10222,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -10685,10 +10333,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -10733,9 +10378,7 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -10754,12 +10397,14 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + assert collected == [ + _anthropic_messages_message_start_chunk(), + _anthropic_messages_content_chunk("hi"), + message_stop, + ] @pytest.mark.asyncio @@ -10771,9 +10416,7 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -10844,7 +10487,9 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source(), initial_kwargs={"model": "primary"} + ) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -12645,9 +12290,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) + assert litellm.Router._declared_param_allowlist( + {"allowed_openai_params": ["reasoning_effort", 3]} + ) == frozenset({"reasoning_effort"}) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -12727,7 +12372,11 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], + [ + {"model_name": "x"}, + {"model_name": "x", "litellm_params": {}}, + {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, + ], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -12956,9 +12605,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" @@ -12991,9 +12638,7 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) + response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) assert response.choices[0].message.content == "from backup-a" diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts index b917fcedaa2..ea6eaf99106 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -80,9 +80,20 @@ describe("hydrateAutoRouterCompression", () => { expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" }); }); - it("treats a missing model key as same-as-routing", () => { + it("treats a missing model key as no model-hop compression, not same-as-routing", () => { const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" }); - expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined }); + expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "none" }); + }); + + it("re-saving a routing-only config leaves the model hop uncompressed", () => { + // Regression: the backend reads an absent model key as no model-hop compression. + // Hydrating it as same-as-routing made opening the router and saving any unrelated + // edit write the routing guardrail onto the model hop, so the model call silently + // started receiving compressed messages. + const stored = { auto_router_routing_compression: "headroom-a" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored)); + expect(rebuilt.auto_router_model_compression).toBe("none"); + expect(rebuilt.auto_router_model_compression).not.toBe("headroom-a"); }); it("round-trips through buildAutoRouterCompressionParams", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index c86416b507f..47d0e3db7e4 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -53,7 +53,11 @@ export const hydrateAutoRouterCompression = (litellmParams: { const routing = litellmParams.auto_router_routing_compression ?? undefined; if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; - const model = litellmParams.auto_router_model_compression ?? undefined; - const sameAsRouting = model === undefined || model === routing; + // An absent model key is no model-hop compression, not same-as-routing: the backend + // reads it as None (policy_from_litellm_params). Hydrating it as same-as-routing + // would make re-saving an unrelated edit write the routing guardrail onto the model + // hop and silently start compressing the model call. + const model = litellmParams.auto_router_model_compression ?? NO_COMPRESSION; + const sameAsRouting = model === routing; return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; }; From 0b3687ec56153225d7b8f2a0c2652bf2f589ce2e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:49:00 -0700 Subject: [PATCH 31/61] fix(shadow_eval): import Final for the test helper's annotation --- tests/test_litellm/integrations/test_shadow_eval_logger.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/integrations/test_shadow_eval_logger.py b/tests/test_litellm/integrations/test_shadow_eval_logger.py index 371f6f75a05..eecd876219e 100644 --- a/tests/test_litellm/integrations/test_shadow_eval_logger.py +++ b/tests/test_litellm/integrations/test_shadow_eval_logger.py @@ -3,6 +3,7 @@ the detached pipeline's single attempt-row write, and the cache-first job lookup import asyncio from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest From d0d09e53438d51b25cb0e0f8a29a329e8d93a7e9 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 09:51:23 -0700 Subject: [PATCH 32/61] feat(router): meter auto-router tier and prompt customization against the auto_router license feature (#39674) Generalizes the heuristic_v2 ceiling from #39468 into a capability table whose records own their in-process predicate, SQL spelling and refusal wording. The existing heuristic_v2 capability keeps its own one-router ceiling. A single customization capability combines operator-defined tier definitions with every operator-written part of the classifier prompt. The prompt half only applies to classifier types that call an LLM. The shipped default prompt, classification rubric presets, tier-label renames and tier model choices remain ungated. Scope every enforcement point to actual complexity routers. A model-less PATCH or legacy update now decrypts the stored model before accepting strategy-router settings, so a regular model cannot acquire a router config or spend a license slot. Under the existing advisory lock, the cross-pod candidate query returns only model scalars and the count decrypts and classifies them in process; old non-router rows carrying a capability-shaped config no longer block a real complexity router. The signed auto_router license feature makes both ceilings unlimited. --- litellm/constants.py | 2 +- litellm/proxy/auth/litellm_license.py | 11 +- .../model_management_endpoints.py | 157 +++++++--- litellm/proxy/proxy_server.py | 34 +- litellm/router.py | 45 +-- .../router_utils/auto_router_model_naming.py | 134 +++++++- litellm/types/router.py | 4 +- .../proxy/auth/test_litellm_license.py | 18 +- .../test_model_management_endpoints.py | 292 +++++++++++++++--- .../proxy/proxy_server/test_proxy_config.py | 91 +++++- .../router_strategy/test_complexity_router.py | 232 +++++++++++++- .../test_auto_router_model_naming.py | 172 +++++++++-- 12 files changed, 987 insertions(+), 205 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index da731cb5eb2..7d6de612349 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -40,7 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", - "heuristic_v2_router_limit", + "auto_router_capability_limit", } ) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 55bb1e3925a..067ac7905c5 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" -HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." +AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." class LicenseCheck: @@ -153,11 +153,12 @@ class LicenseCheck: return False return team_count > _max_teams_in_license - def heuristic_v2_router_limit(self) -> int | None: + def auto_router_capability_limit(self) -> int | None: """ - How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the - signed license lists the auto_router feature, otherwise one. A license verified through - the API carries no feature list, so it does not lift the limit either. + How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined + tier_definitions): unlimited (None) only when the signed license lists the auto_router + feature, otherwise one per capability. A license verified through the API carries no + feature list, so it does not lift the limit either. """ if self.airgapped_license_data is None: return 1 diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index d4e03a05c52..b77108911aa 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -50,7 +50,7 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UserAPIKeyAuth, ) -from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY +from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -98,11 +98,13 @@ from litellm.router_strategy.complexity_router import ( normalize_classification_prompt, ) from litellm.router_utils.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, + capability_limit_violation, carries_complexity_router_settings, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, - uses_heuristic_v2_classifier, + count_capability_routers, + gated_capability_of, + is_complexity_router_model, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -237,11 +239,13 @@ def _strategy_router_write_violation( An auto-router deployment's ``litellm_params.model`` (``auto_router/...``) is the discriminator the router loads it by; a write that mangles it makes the router drop the deployment silently under ``ignore_invalid_deployments``. - Only writes that supply ``litellm_params.model`` are judged on the naming - contract, against the merged (stored + incoming) params, so partial patches - and restores of an already-corrupted row stay legal. A config is judged only - when the write carries one, for the same reason: a rename must not be held - hostage by a stored config it does not touch. Returns the violation, or None. + A patch adding auto-router settings is judged against the effective model, + decrypting the stored model when the patch omits it, so a regular deployment + cannot claim a strategy-router configuration. Unrelated partial patches and + restores that do not touch strategy-router settings stay legal. A config is + judged only when the write carries one, for the same reason: a rename must + not be held hostage by a stored config it does not touch. Returns the + violation, or None. """ if incoming_params is None: return None @@ -256,14 +260,18 @@ def _strategy_router_write_violation( for source in (incoming_params, existing_params) if source is not None and getattr(source, field, None) is not None ) - # Scope reads the incoming model because the stored one is encrypted at rest. - if carries_complexity_router_settings(incoming_params.model, present_fields): + effective_params: Final = _effective_complexity_router_params(incoming_params, existing_params) + effective_model: Final = effective_params.get("model") + if carries_complexity_router_settings( + effective_model if isinstance(effective_model, str) else None, present_fields + ): placement_violation: Final = validate_complexity_router_config_placement(incoming_params.model_extra) if placement_violation is not None: return placement_violation - if incoming_params.model is None: - return None - return validate_strategy_router_model_write(model=incoming_params.model, present_fields=present_fields) + return validate_strategy_router_model_write( + model=effective_model if isinstance(effective_model, str) else "", + present_fields=present_fields, + ) def _raise_on_strategy_router_write_violation( @@ -281,14 +289,23 @@ def _raise_on_strategy_router_write_violation( ) -HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301 -_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" -_HEURISTIC_V2_DB_ROWS_SQL: Final = """ -SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable" +AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 +_CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_STORED_LITELLM_PARAMS_SQL: Final = ( + "(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)" +) +_STORED_COMPLEXITY_CONFIG_SQL: Final = f"{_STORED_LITELLM_PARAMS_SQL} -> 'complexity_router_config'" +_CAPABILITY_DB_ROWS_SQL: Final[Mapping[str, str]] = MappingProxyType( + { + capability.key: f""" +SELECT {_STORED_LITELLM_PARAMS_SQL} ->> 'model' AS model +FROM "LiteLLM_ProxyModelTable" WHERE model_id <> $1 - AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END) - -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2' + AND ({capability.sql_config_predicate.format(config=_STORED_COMPLEXITY_CONFIG_SQL)}) """ + for capability in GATED_AUTO_ROUTER_CAPABILITIES + } +) def _effective_complexity_router_config( @@ -301,13 +318,44 @@ def _effective_complexity_router_config( return existing_params.complexity_router_config -@asynccontextmanager -async def _heuristic_v2_slot( - prisma_client: PrismaClient, *, effective_config: object, model_id: str | None -) -> AsyncGenerator[_ProxyModelTable, None]: - """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled. +def _effective_model( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> str | None: + """The model a write leaves on the row, decrypting an existing value only when the patch omits it.""" + incoming: Final = None if incoming_params is None else incoming_params.model + if incoming is not None: + return incoming + existing: Final = None if existing_params is None else existing_params.model + if existing is None: + return None + decrypted: Final = decrypt_value_helper( + value=existing, + key="model", + exception_type="debug", + return_original_value=True, + ) + return decrypted if isinstance(decrypted, str) else None - A write that leaves the row on classifier_type heuristic_v2 under a limited license runs + +def _effective_complexity_router_params( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> Mapping[str, object]: + """The model and complexity config a write leaves, for placement and capability decisions.""" + return MappingProxyType( + { + "model": _effective_model(incoming_params, existing_params), + "complexity_router_config": _effective_complexity_router_config(incoming_params, existing_params), + } + ) + + +@asynccontextmanager +async def _auto_router_capability_slot( + prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None +) -> AsyncGenerator[_ProxyModelTable, None]: + """Hand out the model table to write through while the row's claim on a licensed capability is settled. + + A write that leaves the row claiming a licensed capability under a limited license runs inside one transaction that takes an advisory lock in its own statement before counting (a statement's snapshot predates anything it locks), so pods cannot both pass the count: the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged @@ -321,21 +369,37 @@ async def _heuristic_v2_slot( """ from litellm.proxy.proxy_server import _license_check, llm_router - limit: Final = _license_check.heuristic_v2_router_limit() - if limit is None or not uses_heuristic_v2_classifier(effective_config): + limit: Final = _license_check.auto_router_capability_limit() + capability: Final = gated_capability_of(effective_params) + if limit is None or capability is None: yield _proxy_model_table(prisma_client) return async with prisma_client.db.tx() as tx_ctx: tables: Final[_TxModelTables] = tx_ctx - await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY) - rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "") - db_held: Final = rows[0].get("held") if rows else 0 + await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( + _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" + ) + db_held: Final = sum( + 1 + for row in rows + for stored_model in (row.get("model"),) + if isinstance(stored_model, str) + and is_complexity_router_model( + decrypt_value_helper( + value=stored_model, + key="model", + exception_type="debug", + return_original_value=True, + ) + ) + ) config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) - held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows) - violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit) + held: Final = db_held + count_capability_routers(config_rows, capability=capability) + violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit) if violation is not None: raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}" + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" ) yield tables.litellm_proxymodeltable await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") @@ -791,6 +855,9 @@ async def patch_model( existing_params=db_model.litellm_params, ) + effective_params: Final = _effective_complexity_router_params( + patch_data.litellm_params, db_model.litellm_params + ) requested_model_name: Final = patch_data.model_name stored_model_name: str | None = None @@ -799,11 +866,9 @@ async def patch_model( stored_model_name = update_data.get("model_name") update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name update_data["updated_at"] = cast(str, get_utc_datetime()) - async with _heuristic_v2_slot( + async with _auto_router_capability_slot( prisma_client, - effective_config=_effective_complexity_router_config( - patch_data.litellm_params, db_model.litellm_params - ), + effective_params=effective_params, model_id=model_id, ) as table: return await table.update(where={"model_id": model_id}, data=update_data) @@ -1959,9 +2024,12 @@ async def add_new_model( model_params=priced_model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, - slot=_heuristic_v2_slot( + slot=_auto_router_capability_slot( prisma_client, - effective_config=priced_model_params.litellm_params.complexity_router_config, + effective_params=_effective_complexity_router_params( + priced_model_params.litellm_params, + None, + ), model_id=priced_model_params.model_info.id, ), ) @@ -2110,6 +2178,9 @@ async def update_model( incoming_params=model_params.litellm_params, existing_params=deployment.litellm_params, ) + effective_params: Final = _effective_complexity_router_params( + model_params.litellm_params, deployment.litellm_params + ) # update DB if store_model_in_db is True: @@ -2147,11 +2218,9 @@ async def update_model( "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, **({} if renamed_to is None else {"model_name": renamed_to}), } - async with _heuristic_v2_slot( + async with _auto_router_capability_slot( prisma_client, - effective_config=_effective_complexity_router_config( - model_params.litellm_params, deployment.litellm_params - ), + effective_params=effective_params, model_id=_model_id, ) as table: model_response: Final = await table.update( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0f0abd0eeae..88e3f79ca52 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -118,10 +118,11 @@ from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, ) from litellm.router_utils.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, + capability_limit_violation, carries_complexity_router_settings, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, + count_capability_routers, validate_complexity_router_config_placement, ) from litellm.types.utils import ( @@ -303,7 +304,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck +from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -4340,17 +4341,28 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object]) raise ValueError(f"model {model.get('model_name', '')!r}: {violation}") -def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: +def validate_auto_router_capability_limits(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: """ - Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows. + Refuse to start when config.yaml defines more auto-routers claiming a licensed capability than allowed. Checked here rather than left to router registration for the same reason as the two validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so the router's own refusal would turn the extra router into a silently missing model. """ - violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit) - if violation is not None: - raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}") + violations: Final = tuple( + message + for capability in GATED_AUTO_ROUTER_CAPABILITIES + if ( + message := capability_limit_violation( + capability=capability, + held=count_capability_routers(model_list, capability=capability), + limit=limit, + ) + ) + is not None + ) + if violations: + raise ValueError(f"config.yaml model_list: {' '.join(violations)} {AUTO_ROUTER_LICENSE_REMEDY}") def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place @@ -5758,7 +5770,7 @@ class ProxyConfig: model_list: Final = config.get("model_list", None) if model_list: router_params["model_list"] = model_list - validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit()) + validate_auto_router_capability_limits(model_list, limit=_license_check.auto_router_capability_limit()) print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" ) @@ -5848,7 +5860,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, - heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, + auto_router_capability_limit=_license_check.auto_router_capability_limit, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6309,7 +6321,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, - heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, + auto_router_capability_limit=_license_check.auto_router_capability_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: diff --git a/litellm/router.py b/litellm/router.py index 6c7611c6236..6943eece90f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -116,10 +116,11 @@ from litellm.router_utils.add_retry_fallback_headers import ( ) from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capability, classify_strategy_router_model, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, - uses_heuristic_v2_classifier, + count_capability_routers, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -208,6 +209,7 @@ from litellm.types.router import ( AlertingConfig, AllowedFailsPolicy, AssistantsTypedDict, + AutoRouterCapabilityLimit, ConsumedRequestTagsStamp, CredentialLiteLLMParams, CustomRoutingStrategyBase, @@ -215,7 +217,6 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, - HeuristicV2RouterLimit, LiteLLM_Params, MockRouterTestingParams, ModelGroupInfo, @@ -692,7 +693,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, - heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None, + auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -769,7 +770,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments - self.heuristic_v2_router_limit = heuristic_v2_router_limit + self.auto_router_capability_limit = auto_router_capability_limit self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks @@ -8811,20 +8812,21 @@ class Router: if not (isinstance(model_info, Mapping) and model_info.get("db_model")): yield deployment - def heuristic_v2_router_limit_violation(self) -> str | None: + def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: """ - Why one more heuristic_v2 router cannot join this router, or None when it can. + Why one more router claiming ``capability`` cannot join this router, or None when it can. Judged against every deployment currently on the model_list; an upsert pops the row being - edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is - resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which - is the SDK default, and the proxy injects a resolver backed by its license. + edited first, so an edit of an existing gated router keeps its own slot. The limit is + resolved on every call through ``auto_router_capability_limit``; unset means unlimited, + which is the SDK default, and the proxy injects a resolver backed by its license. """ - limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None - others: Final = count_heuristic_v2_routers( - deployment for deployment in self.model_list if isinstance(deployment, Mapping) + limit: Final = self.auto_router_capability_limit() if self.auto_router_capability_limit is not None else None + others: Final = count_capability_routers( + (deployment for deployment in self.model_list if isinstance(deployment, Mapping)), + capability=capability, ) - return heuristic_v2_limit_violation(held=others + 1, limit=limit) + return capability_limit_violation(capability=capability, held=others + 1, limit=limit) def init_complexity_router_deployment(self, deployment: Deployment): """ @@ -8843,8 +8845,9 @@ class Router: ) complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config - if uses_heuristic_v2_classifier(complexity_router_config): - limit_violation: Final = self.heuristic_v2_router_limit_violation() + capability: Final = claimed_capability(complexity_router_config) + if capability is not None: + limit_violation: Final = self.auto_router_capability_violation(capability) if limit_violation is not None: raise ValueError(limit_violation) @@ -9674,13 +9677,13 @@ class Router: """Put a deployment back the way it was before a failed upsert popped it. A rollback re-admits state that was already serving, so it does not go through the - heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first + capability ceiling a newcomer gets: with the ceiling tightened since the deployment first registered, judging the rollback would drop a serving router over an unrelated failed edit. """ if previous_deployment is None or self.has_model_id(model_id): return - limit_resolver: Final = self.heuristic_v2_router_limit - self.heuristic_v2_router_limit = None + limit_resolver: Final = self.auto_router_capability_limit + self.auto_router_capability_limit = None try: self.add_deployment(deployment=previous_deployment) verbose_router_logger.info( @@ -9696,7 +9699,7 @@ class Router: restore_error, ) finally: - self.heuristic_v2_router_limit = limit_resolver + self.auto_router_capability_limit = limit_resolver @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index 2efbfb5782e..190c4921d5f 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under ``ignore_invalid_deployments``. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -81,6 +81,11 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: return "semantic" +def is_complexity_router_model(model: str | None) -> bool: + """Whether ``model`` selects the complexity-router implementation.""" + return classify_strategy_router_model(model or "") == "complexity" + + def _named(value: object, role: StrategyRouterDependencyRole) -> tuple[StrategyRouterDependency, ...]: """One dependency from a scalar field, or none when it is absent or not a name.""" return (StrategyRouterDependency(value, role),) if isinstance(value, str) and value else () @@ -168,20 +173,121 @@ def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool: return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2" -def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool: - """Whether this deployment is a complexity router that classifies with heuristic_v2.""" - return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and ( - uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config")) +def defines_custom_tiers(complexity_router_config: object) -> bool: + """Whether this complexity config replaces the built-in tier ladder with operator-defined tier_definitions. + + Mirrors the SQL spelling on the capability record: only an actual array claims the capability, + so an explicit JSON null or a malformed value does not. + """ + return isinstance(_mapping(complexity_router_config).get("tier_definitions"), (list, tuple)) + + +OPERATOR_CLASSIFIER_PROMPT_FIELDS: Final = ("classification_prompt", "classification_examples") + + +def defines_custom_classifier_prompt(complexity_router_config: object) -> bool: + """Whether an operator wrote any part of this router's classifier prompt themselves. + + Three spellings, all metered: a whole replacement prompt (``classifier_llm_config.system_prompt``), + replacement opening instructions (``classification_prompt``), and replacement calibration examples + (``classification_examples``). Choosing a shipped ``classification_rubric`` preset is not authoring. + Scoped to the classifier types that actually call an LLM, which is also where the config validator + accepts these fields: the heuristic scorers never read them. + """ + config: Final = _mapping(complexity_router_config) + if config.get("classifier_type") not in LLM_CLASSIFIER_TYPES: + return False + return _mapping(config.get("classifier_llm_config")).get("system_prompt") is not None or any( + config.get(field) is not None for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS ) -def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int: - """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers.""" - return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params")))) +def uses_custom_tier_or_classifier_prompt(complexity_router_config: object) -> bool: + """Whether this router replaces shipped tiers or its shipped classifier prompt.""" + return defines_custom_tiers(complexity_router_config) or defines_custom_classifier_prompt(complexity_router_config) -def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: - """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits. +_LLM_CLASSIFIER_TYPES_SQL: Final = ", ".join(f"'{name}'" for name in sorted(LLM_CLASSIFIER_TYPES)) + + +@dataclass(frozen=True, slots=True) +class GatedAutoRouterCapability: + """A complexity-router capability the license meters, in every spelling an enforcement point needs. + + ``uses`` and ``sql_config_predicate`` answer the same question, in process and in a DB count over + stored ``litellm_params`` (``{config}`` is the caller's expression for the normalized + ``complexity_router_config`` jsonb, substituted as many times as the predicate needs); they live + on one record so they cannot drift apart. ``subject`` and ``remedy`` build the shared refusal + message. A validated config claims at most one capability, and the validator is what makes that + true: tier_definitions rejects every heuristic classifier_type, and it also rejects the + classifier system_prompt, which in turn only applies to the classifier types heuristic_v2 is not. + """ + + key: str + subject: str + remedy: str + uses: Callable[[object], bool] + sql_config_predicate: str + + +HEURISTIC_V2_CAPABILITY: Final = GatedAutoRouterCapability( + key="heuristic_v2", + subject="with classifier_type 'heuristic_v2'", + remedy="Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router.", + uses=uses_heuristic_v2_classifier, + sql_config_predicate="{config} ->> 'classifier_type' = 'heuristic_v2'", +) + +_OPERATOR_PROMPT_FIELDS_SQL: Final = " OR ".join( + f"{{config}} ->> '{field}' IS NOT NULL" for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS +) + +CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability( + key="tier_or_classifier_prompt", + subject="with operator-defined tier_definitions or an operator-written classifier prompt", + remedy=( + "Use the shipped tiers and classifier prompt for this router or remove an existing router " + "with tier_definitions or its own classifier prompt." + ), + uses=uses_custom_tier_or_classifier_prompt, + sql_config_predicate=( + "jsonb_typeof({config} -> 'tier_definitions') = 'array' OR " + f"({{config}} ->> 'classifier_type' IN ({_LLM_CLASSIFIER_TYPES_SQL}) AND (" + "{config} -> 'classifier_llm_config' ->> 'system_prompt' IS NOT NULL OR " + f"{_OPERATOR_PROMPT_FIELDS_SQL}))" + ), +) + +GATED_AUTO_ROUTER_CAPABILITIES: Final = (HEURISTIC_V2_CAPABILITY, CUSTOMIZATION_CAPABILITY) + + +def claimed_capability(complexity_router_config: object) -> GatedAutoRouterCapability | None: + """The licensed capability this complexity config claims, or None.""" + return next( + (capability for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(complexity_router_config)), + None, + ) + + +def gated_capability_of(litellm_params: Mapping[str, object]) -> GatedAutoRouterCapability | None: + """The licensed capability this deployment claims, or None unless it is a complexity router.""" + model: Final = litellm_params.get("model") + if not is_complexity_router_model(model if isinstance(model, str) else None): + return None + return claimed_capability(litellm_params.get("complexity_router_config")) + + +def count_capability_routers( + deployments: Iterable[Mapping[str, object]], *, capability: GatedAutoRouterCapability +) -> int: + """How many of ``deployments`` (router model_list entries or config.yaml rows) claim ``capability``.""" + return sum( + 1 for deployment in deployments if gated_capability_of(_mapping(deployment.get("litellm_params"))) is capability + ) + + +def capability_limit_violation(*, capability: GatedAutoRouterCapability, held: int, limit: int | None) -> str | None: + """Why holding ``held`` routers claiming ``capability`` exceeds ``limit``, or None when it fits. ``limit`` None means unlimited. The message is shared by every enforcement point (config load, model writes, router registration) and stays SDK-neutral: it names the cap and what @@ -190,8 +296,8 @@ def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: if limit is None or held <= limit: return None return ( - f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make " - f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router." + f"At most {limit} auto-router(s) {capability.subject} can be registered but this would make " + f"{held}. {capability.remedy}" ) @@ -237,9 +343,7 @@ def carries_complexity_router_settings(model: str | None, present_fields: frozen ``validate_strategy_router_model_write`` is judged on, so a router named only by its default model is in scope, and a field added to the table above is covered here for free. """ - return classify_strategy_router_model(model or "") == "complexity" or bool( - present_fields & _COMPLEXITY_ROUTER_FIELDS - ) + return is_complexity_router_model(model) or bool(present_fields & _COMPLEXITY_ROUTER_FIELDS) def validate_complexity_router_config_placement(litellm_params: Mapping[str, object] | None) -> str | None: diff --git a/litellm/types/router.py b/litellm/types/router.py index 267e8853db1..728d1037f3d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -887,9 +887,9 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... -class HeuristicV2RouterLimit(Protocol): +class AutoRouterCapabilityLimit(Protocol): """ - Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited. + Resolves how many complexity routers may claim each licensed capability right now; None means unlimited. The Router calls it on every registration and limit query instead of caching the answer, so the proxy can keep the limit on its license object (re-verified on config load) rather than hand diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py index 1db53638070..d3f80982c7a 100644 --- a/tests/test_litellm/proxy/auth/test_litellm_license.py +++ b/tests/test_litellm/proxy/auth/test_litellm_license.py @@ -34,27 +34,27 @@ def test_is_over_limit(): assert license_check.is_over_limit(99) is False -def test_heuristic_v2_router_limit() -> None: +def test_auto_router_capability_limit() -> None: """Only the signed license's auto_router feature lifts the one-router limit; an API-verified license (no airgapped data) and an airgapped license without the feature keep it.""" license_check = LicenseCheck() license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]} - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None license_check.airgapped_license_data = { "expiration_date": "2999-01-01", "allowed_features": ["sso", "auto_router", "audit_logs"], } - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]} - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 license_check.airgapped_license_data = {"expiration_date": "2999-01-01"} - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 license_check.airgapped_license_data = None - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]: @@ -81,12 +81,12 @@ def test_expired_or_unreadable_license_grants_no_features() -> None: license_check = LicenseCheck() public_key, valid_key = _signed_license("2999-01-01") assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None _, expired_key = _signed_license("2000-01-01") assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True assert license_check.airgapped_license_data is None - assert license_check.heuristic_v2_router_limit() == 1 + assert license_check.auto_router_capability_limit() == 1 assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True @@ -98,4 +98,4 @@ def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None: public_key, license_key = _signed_license("2999-01-01") assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True - assert license_check.heuristic_v2_router_limit() is None + assert license_check.auto_router_capability_limit() is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 3edeeedbae9..33de2a09626 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -2,6 +2,7 @@ import inspect import asyncio import contextlib import json +from collections.abc import Mapping from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -4048,6 +4049,72 @@ class TestStrategyRouterWriteValidation: ) assert _strategy_router_write_violation(incoming_params=None, existing_params=None) is None + @pytest.mark.parametrize( + "config", + [ + {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tier_definitions": [ + {"name": "routine", "description": "routine drafting"}, + {"name": "hard", "description": "hard reasoning"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + }, + ], + ) + def test_model_less_patch_cannot_attach_router_config_to_a_regular_model(self, config: dict[str, object]) -> None: + """The license gate applies only to complexity routers, so a partial PATCH cannot poison a regular + model with a capability-shaped config and make it occupy a slot.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + violation = _strategy_router_write_violation( + incoming_params=updateLiteLLMParams(complexity_router_config=config), + existing_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + ) + + assert violation is not None + assert "does not start with 'auto_router/'" in violation + assert "complexity_router_config" in violation + + def test_effective_params_decrypts_a_stored_complexity_router_model(self, monkeypatch) -> None: + """A database row encrypts model, so the model-aware gate must not accidentally rely on plaintext mocks.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _effective_complexity_router_params, + ) + from litellm.types.router import updateLiteLLMParams + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt") + encrypted_model = encrypt_value_helper("auto_router/complexity_router") + effective_params = _effective_complexity_router_params( + updateLiteLLMParams(complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}), + LiteLLM_Params(model=encrypted_model), + ) + + assert effective_params["model"] == "auto_router/complexity_router" + + def test_model_less_patch_keeps_a_complexity_router_in_scope(self) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _strategy_router_write_violation, + ) + from litellm.types.router import updateLiteLLMParams + + assert ( + _strategy_router_write_violation( + incoming_params=updateLiteLLMParams( + complexity_router_config={"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} + ), + existing_params=self._stored_complexity_params(), + ) + is None + ) + def test_restore_of_corrupted_row_is_allowed(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( _strategy_router_write_violation, @@ -4354,33 +4421,33 @@ class TestStrategyRouterWriteValidation: ) @staticmethod - def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router: + def _live_router_holding_one_capability(limit: int | None, config: Mapping[str, object]) -> Router: return Router( model_list=[ {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}}, { - "model_name": "held-v2", + "model_name": "held", "litellm_params": { "model": "auto_router/complexity_router", - "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + "complexity_router_config": config, }, "model_info": {"id": "held-id"}, }, ], - heuristic_v2_router_limit=lambda: limit, + auto_router_capability_limit=lambda: limit, ) class _FakeTx: - """Stands in for a prisma transaction: records the raw statements and exposes the model table.""" + """Stands in for a prisma transaction: records raw statements and returns encrypted-model candidates.""" - def __init__(self, db_held: int) -> None: - self.db_held = db_held + def __init__(self, db_models: list[str]) -> None: + self.db_models = db_models self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: self.raw_calls.append((sql, args)) - return [{"held": self.db_held}] if "count(*)" in sql else [] + return [{"model": model} for model in self.db_models] if "AS model" in sql else [] async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": return self @@ -4391,9 +4458,9 @@ class TestStrategyRouterWriteValidation: class _FakeDb: """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" - def __init__(self, db_held: int, existing_row: object = None) -> None: + def __init__(self, db_models: list[str], existing_row: object = None) -> None: self.db = self - self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held) + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models) self.litellm_proxymodeltable = MagicMock( create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) ) @@ -4403,6 +4470,43 @@ class TestStrategyRouterWriteValidation: _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}} + _CUSTOM_TIERS = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tier_definitions": [ + {"name": "routine", "description": "routine drafting"}, + {"name": "hard", "description": "hard reasoning"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + } + _TIER_LABELS_ONLY = { + "classifier_type": "heuristic", + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "tier_labels": {"SIMPLE": "Cheap"}, + } + _CUSTOM_PROMPT = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + } + _OPERATOR_EXAMPLES = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "classification_examples": '- "reset my password" -> SIMPLE', + } + _OPERATOR_OPENING_PROMPT = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "classification_prompt": "Grade by data sensitivity", + } + _SHIPPED_RUBRIC = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "classification_rubric": "agentic"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + } @pytest.mark.parametrize( "incoming,existing,expected", @@ -4431,41 +4535,55 @@ class TestStrategyRouterWriteValidation: @pytest.mark.asyncio @pytest.mark.parametrize( - "limit,effective_config,db_held,config_holds_one,model_id,expected", + "limit,effective_params,db_models,config_config,model_id,expected", [ - (1, _V2, 1, False, None, "refused"), - (1, _V2, 0, True, None, "refused"), - (1, _V2, 0, False, None, "reserved"), - (1, _V2, 0, False, "held-id", "reserved"), - (2, _V2, 1, False, None, "reserved"), - (1, _V1, 5, True, None, "plain"), - (1, None, 5, True, None, "plain"), - (None, _V2, 5, True, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], _V2, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, None, "reserved"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, "held-id", "reserved"), + (2, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "reserved"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, ["openai/gpt-4o"], None, None, "reserved"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, [], _CUSTOM_PROMPT, None, "refused"), + (1, {"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIERS}, ["auto_router/complexity_router"], None, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V1}, ["auto_router/complexity_router"], _V2, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": None}, ["auto_router/complexity_router"], _V2, None, "plain"), + (None, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], _V2, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _TIER_LABELS_ONLY}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_PROMPT}, ["auto_router/complexity_router"], None, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, [], _CUSTOM_TIERS, None, "refused"), + (1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_OPENING_PROMPT}, [], _CUSTOM_PROMPT, None, "refused"), + (None, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"), ], ) - async def test_heuristic_v2_slot_matrix( + async def test_auto_router_capability_slot_matrix( self, limit: int | None, - effective_config: object, - db_held: int, - config_holds_one: bool, + effective_params: Mapping[str, object], + db_models: list[str], + config_config: Mapping[str, object] | None, model_id: str | None, expected: str, ) -> None: - """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows - (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL - parameter, and every other write runs on the plain client with no lock.""" + """The slot is claimed inside a locked transaction only for a write that claims a licensed capability + under a limit; the DB rows (other pods included) plus config.yaml routers decide, the row being edited + is excluded through the SQL parameter, and every other write runs on the plain client with no lock. + + heuristic_v2 has its own slot, while custom tier definitions and custom prompts count into one shared + customization slot. Renaming built-in tiers through tier_labels claims nothing at all.""" from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import ( - HEURISTIC_V2_SLOT_LOCK_KEY, - _heuristic_v2_slot, + AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY, + _auto_router_capability_slot, ) + from litellm.router_utils.auto_router_model_naming import gated_capability_of - fake = self._FakeDb(db_held) - live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None + capability = gated_capability_of(effective_params) + + fake = self._FakeDb(db_models) + live_router = self._live_router_holding_one_capability(limit, config_config) if config_config is not None else None with ( - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", @@ -4474,13 +4592,15 @@ class TestStrategyRouterWriteValidation: ): if expected == "refused": with pytest.raises(HTTPException) as exc_info: - async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id): + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id): pass assert exc_info.value.status_code == 403 + assert capability is not None assert "At most 1 auto-router" in str(exc_info.value.detail) + assert capability.subject in str(exc_info.value.detail) assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) return - async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables: + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id) as tables: handle = tables if expected == "plain": await handle.create(data={}) @@ -4489,10 +4609,13 @@ class TestStrategyRouterWriteValidation: return assert handle is fake.tx_obj.litellm_proxymodeltable published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") - (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls + (lock_sql, lock_params), (count_sql, count_params) = fake.tx_obj.raw_calls assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql - assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,) + assert lock_params == (AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,) assert count_params == (model_id or "",) + assert "AS model" in count_sql + assert capability is not None + assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql @pytest.mark.asyncio async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: @@ -4549,14 +4672,14 @@ class TestStrategyRouterWriteValidation: ) admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - fake = self._FakeDb(db_held=1) + fake = self._FakeDb(["auto_router/complexity_router"]) with ( patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -4579,6 +4702,93 @@ class TestStrategyRouterWriteValidation: fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() fake.litellm_proxymodeltable.create.assert_not_awaited() + @pytest.mark.asyncio + async def test_model_less_patch_rejects_router_config_on_a_regular_model(self) -> None: + """PATCH rejects the poison before its row write or the capability slot.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + from litellm.types.router import updateLiteLLMParams + + model_id = "regular-model" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + regular = Deployment( + model_name="regular-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info={"id": model_id}, + ) + fake = self._FakeDb([]) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag + patch( # test-quality-ok: inject stored regular row without a database + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=regular), + ), + patch( # test-quality-ok: endpoint must reject before database authorization needs a live store + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS) + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.code == "400" + assert "does not start with 'auto_router/'" in str(exc_info.value.message) + assert fake.tx_obj.raw_calls == [] + assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0 + assert fake.litellm_proxymodeltable.update.await_count == 0 + + @pytest.mark.asyncio + async def test_model_less_legacy_update_rejects_router_config_on_a_regular_model(self) -> None: + """The legacy update endpoint enforces the same boundary before its row write or slot.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + from litellm.types.router import ModelInfo, updateLiteLLMParams + + model_id = "regular-model" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + regular = Deployment( + model_name="regular-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), + model_info={"id": model_id}, + ) + existing_row = MagicMock() + existing_row.model_dump.return_value = regular.model_dump() + existing_row.litellm_params = regular.litellm_params.model_dump() + fake = self._FakeDb([], existing_row=existing_row) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag + patch( # test-quality-ok: endpoint must reject before database authorization needs a live store + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.code == "400" + assert "does not start with 'auto_router/'" in str(exc_info.value.message) + assert fake.tx_obj.raw_calls == [] + assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0 + assert fake.litellm_proxymodeltable.update.await_count == 0 + @pytest.mark.asyncio async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" @@ -4591,14 +4801,14 @@ class TestStrategyRouterWriteValidation: model_id = "other-id" admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - fake = self._FakeDb(db_held=1) + fake = self._FakeDb(["auto_router/complexity_router"]) with ( patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: the write must be refused before this DB step runs "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=self._db_complexity_router(model_id)), @@ -4643,14 +4853,14 @@ class TestStrategyRouterWriteValidation: "model_info": {"id": model_id}, } existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"] - fake = self._FakeDb(db_held=1, existing_row=existing_row) + fake = self._FakeDb(["auto_router/complexity_router"], existing_row=existing_row) with ( patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index dcfad8f6815..2babfe432f3 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,7 +28,7 @@ from litellm.proxy.proxy_server import ( resolve_routing_plugins, validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, - validate_heuristic_v2_router_limit, + validate_auto_router_capability_limits, ) from .conftest import normalize @@ -204,13 +204,71 @@ def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> } -def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None: +def _custom_tier_row(model_name: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "tier_definitions": [ + {"name": "routine", "description": "routine drafting"}, + {"name": "hard", "description": "hard reasoning"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + }, + }, + } + + +def _operator_examples_row(model_name: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + "classification_examples": '- "reset my password" -> SIMPLE', + }, + }, + } + + +def _custom_prompt_row(model_name: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, + "tiers": {"SIMPLE": "gpt-4o-mini"}, + }, + }, + } + + +@pytest.mark.parametrize( + "over_limit_rows,subject", + [ + ([_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], "heuristic_v2"), + ([_custom_tier_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "tier_definitions"), + ([_custom_prompt_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"), + ([_custom_tier_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"), + ([_operator_examples_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"), + ], +) +def test_validate_auto_router_capability_limits_refuses_to_start_over_the_limit( + over_limit_rows: list[dict[str, object]], subject: str +) -> None: """Same reason as the two validators above: the proxy router swallows registration errors, so an over-limit config.yaml must fail here instead of booting with a silently missing router.""" with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info: - validate_heuristic_v2_router_limit( - [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1 - ) + validate_auto_router_capability_limits(over_limit_rows, limit=1) + assert subject in str(exc_info.value) assert "'auto_router' feature lifts the limit" in str(exc_info.value) @@ -220,12 +278,15 @@ def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None), ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1), ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1), + ([_custom_tier_row("a"), _custom_tier_row("b")], None), + ([_custom_tier_row("a"), _heuristic_v2_row("b")], 1), ], ) -def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone( +def test_validate_auto_router_capability_limits_leaves_configs_within_the_limit_alone( model_list: list[dict[str, object]], limit: int | None ) -> None: - assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None + """The last case is the separate-ceiling invariant: one router of each capability fits under a limit of one.""" + assert validate_auto_router_capability_limits(model_list, limit=limit) is None _TWO_HEURISTIC_V2_ROUTERS_YAML = ( @@ -247,7 +308,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = ( " classifier_type: heuristic_v2\n" " tiers: {SIMPLE: gpt-4o-mini}\n" "router_settings:\n" - " heuristic_v2_router_limit: 99\n" + " auto_router_capability_limit: 99\n" ) @@ -256,7 +317,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = ( async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only( tmp_path, monkeypatch, license_limit: int | None ) -> None: - """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator + """`router_settings.auto_router_capability_limit` is managed outside config.yaml: an operator cannot grant the entitlement by editing the config, and a licensed proxy boots both routers.""" f = tmp_path / "c.yaml" f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML) @@ -264,15 +325,15 @@ async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_lic monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) monkeypatch.setattr( - "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: license_limit ) if license_limit is None: router, _model_list, _general_settings = await ProxyConfig().load_config( router=None, config_file_path=str(f) ) - assert router.heuristic_v2_router_limit is not None - assert router.heuristic_v2_router_limit() is None + assert router.auto_router_capability_limit is not None + assert router.auto_router_capability_limit() is None assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] return @@ -296,12 +357,12 @@ async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_b monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) - monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1) router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f)) - assert router.heuristic_v2_router_limit is not None - assert router.heuristic_v2_router_limit() == 1 + assert router.auto_router_capability_limit is not None + assert router.auto_router_capability_limit() == 1 assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"}) assert router.upsert_deployment(db_row) is None diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 52e58304476..918ec7bc100 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -15,6 +15,12 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.router_utils.auto_router_model_naming import ( + CUSTOMIZATION_CAPABILITY, + GATED_AUTO_ROUTER_CAPABILITIES, + HEURISTIC_V2_CAPABILITY, + count_capability_routers, +) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY @@ -46,7 +52,6 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, TrainedTierArtifact, ) -from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers from litellm.types.router import ( Deployment, LiteLLM_Params, @@ -1124,7 +1129,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-b", "id-b", "heuristic_v2"), self._router_row("v1-c", "id-c", "heuristic"), ], - heuristic_v2_router_limit=lambda: 1, + auto_router_capability_limit=lambda: 1, ignore_invalid_deployments=True, ) @@ -1139,7 +1144,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: 1, + auto_router_capability_limit=lambda: 1, ) def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None: @@ -1152,14 +1157,14 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: limits["value"], + auto_router_capability_limit=lambda: limits["value"], ignore_invalid_deployments=True, ) assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] - assert router.heuristic_v2_router_limit_violation() is None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None limits["value"] = 1 - assert router.heuristic_v2_router_limit_violation() is not None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] @@ -1174,7 +1179,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: limits["value"], + auto_router_capability_limit=lambda: limits["value"], ignore_invalid_deployments=True, ) limits["value"] = 1 @@ -1195,7 +1200,7 @@ class TestRouterComplexityDeploymentMethods: assert router.upsert_deployment(Deployment(**db_row)) is not None assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"] - assert count_heuristic_v2_routers(router.config_deployments()) == 1 + assert count_capability_routers(router.config_deployments(), capability=HEURISTIC_V2_CAPABILITY) == 1 def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None: """A rollback after a failed upsert re-admits state that was already serving, so it must not be @@ -1208,7 +1213,7 @@ class TestRouterComplexityDeploymentMethods: self._router_row("v2-a", "id-a", "heuristic_v2"), self._router_row("v2-b", "id-b", "heuristic_v2"), ], - heuristic_v2_router_limit=lambda: limits["value"], + auto_router_capability_limit=lambda: limits["value"], ignore_invalid_deployments=True, ) limits["value"] = 1 @@ -1220,7 +1225,7 @@ class TestRouterComplexityDeploymentMethods: assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] live = router.get_deployment(model_id="id-a") assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2" - assert router.heuristic_v2_router_limit_violation() is not None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None: router = Router( @@ -1232,18 +1237,18 @@ class TestRouterComplexityDeploymentMethods: ) assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] - assert router.heuristic_v2_router_limit_violation() is None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None - def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None: + def test_auto_router_capability_violation_frees_the_slot_of_the_router_being_edited(self) -> None: """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot while a different deployment switching to heuristic_v2 is refused.""" router = Router( model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")], - heuristic_v2_router_limit=lambda: 1, + auto_router_capability_limit=lambda: 1, ignore_invalid_deployments=True, ) - assert router.heuristic_v2_router_limit_violation() is not None + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2") assert router.upsert_deployment(Deployment(**edited)) is not None @@ -1254,6 +1259,205 @@ class TestRouterComplexityDeploymentMethods: assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"] + @staticmethod + def _custom_tier_row(model_name: str, model_id: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tier_definitions": [ + {"name": "routine", "description": "routine drafting and lookups"}, + {"name": "hard", "description": "multi-step reasoning under tradeoffs"}, + ], + "tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"}, + "fallback_tier": "routine", + }, + }, + "model_info": {"id": model_id}, + } + + @staticmethod + def _custom_prompt_row(model_name: str, model_id: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + } | {"model_info": {"id": model_id}} + + def test_a_second_custom_prompt_router_is_refused_under_the_ceiling(self) -> None: + """An operator-written classifier system_prompt is metered like the other licensed capabilities.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._custom_prompt_row("prompt-a", "id-a"), + self._custom_prompt_row("prompt-b", "id-b"), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None: + """Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no + prompt at all, leaves a router unmetered, so several of them register under a ceiling of one.""" + def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]: + llm_config: dict[str, object] = {"model": "gpt-4o-mini"} + if preset is not None: + llm_config["classification_rubric"] = preset + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": llm_config, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + "model_info": {"id": model_id}, + } + + router = Router( + model_list=[ + self._POOL, + rubric("default-a", "id-a", None), + rubric("preset-b", "id-b", "agentic"), + rubric("preset-c", "id-c", "chat"), + ], + auto_router_capability_limit=lambda: 1, + ) + + assert sorted(router.complexity_routers) == ["default-a", "preset-b", "preset-c"] + + def test_a_second_custom_tier_router_is_refused_under_the_ceiling(self) -> None: + """Operator-defined tier sets are metered like heuristic_v2: one per proxy without the license.""" + with pytest.raises(ValueError, match="tier_definitions"): + Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._custom_tier_row("tiers-b", "id-b"), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_custom_tier_routers_are_unlimited_with_the_license_feature(self) -> None: + router = Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._custom_tier_row("tiers-b", "id-b"), + ], + auto_router_capability_limit=lambda: None, + ) + + assert sorted(router.complexity_routers) == ["tiers-a", "tiers-b"] + assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is None + + def test_each_capability_holds_its_own_slot(self) -> None: + """heuristic_v2 has its own slot, while custom tiers and custom prompts share one customization + slot: one v2 plus EITHER customization fits, but a second customization of any form is refused.""" + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._custom_tier_row("tiers-a", "id-t"), + ], + auto_router_capability_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"] + assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None + assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is not None + + assert router.upsert_deployment(Deployment(**self._custom_tier_row("tiers-b", "id-t2"))) is None + assert router.upsert_deployment(Deployment(**self._custom_prompt_row("prompt-b", "id-p2"))) is None + assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"] + + @staticmethod + def _operator_prompt_row(model_name: str, model_id: str, field: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_default_model": "gpt-4o-mini", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini"}, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + field: '- "reset my password" -> SIMPLE', + }, + }, + "model_info": {"id": model_id}, + } + + @pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"]) + def test_operator_written_prompt_sections_claim_the_customization_slot(self, field: str) -> None: + """The dashboard prompt editor writes opening instructions and calibration examples as their own + fields on a BUILT-IN tier router, so each must claim the slot on its own.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._operator_prompt_row("prompt-a", "id-a", field), + self._operator_prompt_row("prompt-b", "id-b", field), + ], + auto_router_capability_limit=lambda: 1, + ) + + @pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"]) + def test_an_operator_prompt_section_claims_the_slot_held_by_custom_tiers(self, field: str) -> None: + """Switching the FORM of customization cannot buy a second unlicensed router.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._operator_prompt_row("prompt-b", "id-b", field), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_a_custom_prompt_claims_the_slot_held_by_custom_tiers(self) -> None: + """The customization ceiling is shared: changing its form cannot get a second unlicensed router.""" + with pytest.raises(ValueError, match="operator-written classifier prompt"): + Router( + model_list=[ + self._POOL, + self._custom_tier_row("tiers-a", "id-a"), + self._custom_prompt_row("prompt-b", "id-b"), + ], + auto_router_capability_limit=lambda: 1, + ) + + def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None: + """tier_labels renames the built-in ladder without defining one, so it stays ungated: two such + routers register under a ceiling of one.""" + def labeled(model_name: str, model_id: str) -> dict[str, object]: + row = self._router_row(model_name, model_id, "heuristic") + row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"} + return row + + router = Router( + model_list=[self._POOL, labeled("labels-a", "id-a"), labeled("labels-b", "id-b")], + auto_router_capability_limit=lambda: 1, + ) + + assert sorted(router.complexity_routers) == ["labels-a", "labels-b"] + def test_hybrid_initialization_waits_for_later_pool_deployments(self): router = Router( model_list=[ diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 238d0546518..8dede941a14 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -5,9 +5,11 @@ import pytest from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, - count_heuristic_v2_routers, - heuristic_v2_limit_violation, - is_heuristic_v2_router, + GATED_AUTO_ROUTER_CAPABILITIES, + capability_limit_violation, + claimed_capability, + count_capability_routers, + gated_capability_of, strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, @@ -376,38 +378,122 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie assert carries_complexity_router_settings(model, present_fields) is scoped +_HV2_CONFIG: Mapping[str, object] = {"classifier_type": "heuristic_v2"} +_CUSTOM_TIER_CONFIG: Mapping[str, object] = { + "classifier_type": "llm", + "tier_definitions": [{"name": "routine", "description": "easy"}, {"name": "hard", "description": "hard"}], +} +_CUSTOM_PROMPT_CONFIG: Mapping[str, object] = { + "classifier_type": "llm", + "classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"}, +} + + @pytest.mark.parametrize( - "litellm_params,expected", + "config,expected_key", [ - ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), - ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), - ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False), - ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False), - ({"model": "auto_router/complexity_router"}, False), - ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), - ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), - ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False), - ({}, False), + (_CUSTOM_PROMPT_CONFIG, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": "grade it"}, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_examples": '- "x" -> SIMPLE'}, "tier_or_classifier_prompt"), + ({"classifier_type": "hybrid", "classification_examples": "- y -> MEDIUM"}, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": None, "classification_examples": None}, None), + ({"classifier_type": "heuristic", "classification_examples": "- x -> SIMPLE"}, None), + ({"classifier_type": "hybrid", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"), + ({"classifier_type": "heuristic_first", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "classification_rubric": "chat"}}, None), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}}, None), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": None}}, None), + ({"classifier_type": "heuristic", "classifier_llm_config": {"system_prompt": "p"}}, None), + ({"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}}, "heuristic_v2"), + ({"classifier_type": "llm", "classifier_llm_config": "not a mapping"}, None), ], ) -def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None: - """Only a complexity router whose config selects heuristic_v2 counts toward the license limit.""" - assert is_heuristic_v2_router(litellm_params) is expected +def test_custom_classifier_prompt_capability(config: Mapping[str, object], expected_key: str | None) -> None: + """Every operator-written part of the classifier prompt claims the customization slot: a whole + replacement system_prompt, replacement opening instructions (classification_prompt), or replacement + calibration examples (classification_examples). + + A shipped rubric preset stays free, and the heuristic scorers never read system_prompt, so a + value sitting on one is inert and claims nothing (heuristic_v2 still claims its own capability). + """ + claimed = claimed_capability(config) + assert (None if claimed is None else claimed.key) == expected_key -def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None: - v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}} +@pytest.mark.parametrize( + "model,expected", + [ + ("auto_router/complexity_router", True), + ("auto_router/complexity_router-eu", True), + ("auto_router/semantic_router", False), + ("auto_router/adaptive_router", False), + ("auto_router/quality_router", False), + ("openai/gpt-4o", False), + (None, False), + ], +) +def test_is_complexity_router_model(model: str | None, expected: bool) -> None: + from litellm.router_utils.auto_router_model_naming import is_complexity_router_model + + assert is_complexity_router_model(model) is expected + + +@pytest.mark.parametrize( + "litellm_params,expected_key", + [ + ({"model": "auto_router/complexity_router", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"), + ({"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_definitions": None}}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_labels": {"SIMPLE": "Cheap"}}}, None), + ({"model": "auto_router/complexity_router"}, None), + ({"model": "auto_router/quality_router", "complexity_router_config": _HV2_CONFIG}, None), + ({"model": "auto_router/quality_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None), + ({"model": "openai/gpt-4o", "complexity_router_config": _HV2_CONFIG}, None), + ({"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None), + ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, None), + ({}, None), + ], +) +def test_gated_capability_of(litellm_params: Mapping[str, object], expected_key: str | None) -> None: + """Only a complexity router claiming a licensed capability counts toward that capability's limit. + + Renaming the built-in tiers through tier_labels is not a custom tier set, so it stays ungated. + """ + capability = gated_capability_of(litellm_params) + assert (None if capability is None else capability.key) == expected_key + + +@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key) +def test_count_capability_routers_counts_only_its_own_capability(capability) -> None: + """Each capability has its own ceiling, so a router claiming the sibling capability never counts, + while a custom tier set and a custom classifier prompt count into the SAME customization slot.""" + def row(name: str, config: Mapping[str, object] | None) -> Mapping[str, object]: + params = {"model": "auto_router/complexity_router"} | ({} if config is None else {"complexity_router_config": config}) + return {"model_name": name, "litellm_params": params} + + by_key = { + "heuristic_v2": (_HV2_CONFIG, _HV2_CONFIG), + "tier_or_classifier_prompt": (_CUSTOM_TIER_CONFIG, _CUSTOM_PROMPT_CONFIG), + } + mine_first, mine_second = by_key[capability.key] + theirs = next(configs[0] for key, configs in by_key.items() if key != capability.key) rows: list[Mapping[str, object]] = [ - {"model_name": "a", "litellm_params": v2}, - {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}}, - {"model_name": "c", "litellm_params": v2}, - {"model_name": "d"}, - {"model_name": "e", "litellm_params": "not a mapping"}, + row("a", mine_first), + row("b", theirs), + {"model_name": "c", "litellm_params": {"model": "openai/gpt-4o"}}, + row("d", mine_second), + {"model_name": "e"}, + {"model_name": "f", "litellm_params": "not a mapping"}, ] - assert count_heuristic_v2_routers(rows) == 2 - assert count_heuristic_v2_routers(()) == 0 + assert count_capability_routers(rows, capability=capability) == 2 + assert count_capability_routers((), capability=capability) == 0 +@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key) @pytest.mark.parametrize( "held,limit,violates", [ @@ -419,10 +505,42 @@ def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ (4, 3, True), ], ) -def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None: - violation = heuristic_v2_limit_violation(held=held, limit=limit) +def test_capability_limit_violation(held: int, limit: int | None, violates: bool, capability) -> None: + violation = capability_limit_violation(capability=capability, held=held, limit=limit) assert (violation is not None) is violates if violation is not None: assert f"At most {limit} auto-router" in violation assert f"would make {held}" in violation + assert capability.subject in violation + assert capability.remedy in violation assert "license" not in violation + + +def test_every_gated_capability_has_a_distinct_predicate_and_sql_spelling() -> None: + """The in-process and SQL halves of a capability must stay paired, and no two capabilities may collide.""" + keys = tuple(capability.key for capability in GATED_AUTO_ROUTER_CAPABILITIES) + assert len(set(keys)) == len(keys) + for capability in GATED_AUTO_ROUTER_CAPABILITIES: + assert "{config}" in capability.sql_config_predicate + assert capability.uses is not None + + +@pytest.mark.parametrize( + "config", + [ + _HV2_CONFIG, + _CUSTOM_TIER_CONFIG, + _CUSTOM_PROMPT_CONFIG, + {"classifier_type": "heuristic"}, + {"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}}, + {"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": "p"}, "tier_labels": {"SIMPLE": "Cheap"}}, + ], +) +def test_capabilities_are_mutually_exclusive_on_one_config(config: Mapping[str, object]) -> None: + """No config claims two capabilities, which is what lets one lock and one count serve them all. + + The config validator is what makes this true and is pinned separately in test_complexity_router: + tier_definitions rejects every heuristic classifier_type and rejects the classifier system_prompt, + and system_prompt only counts for the classifier types heuristic_v2 is not one of. + """ + assert sum(1 for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(config)) <= 1 From 8284208af261bd32d78ff3fb43040117894fd358 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 09:51:51 -0700 Subject: [PATCH 33/61] fix(auto-router compression): restrict both hops to real compression guardrails The two policy fields are operator-supplied names and nothing else constrained them. The routing hop calls apply_guardrail directly, which hands the guardrail the conversation and POSTs it to whatever service backs that guardrail, and the model hop is added to metadata["guardrails"], which runs it even when it is not default_on. So naming an ordinary guardrail turned either hop into a way to invoke it and ship prompt content to it. Both hops now refuse a name that does not resolve to an active compression guardrail, and say so in the log rather than failing quietly. --- .../guardrails/auto_router_compression.py | 52 ++++++++++++++--- .../test_auto_router_compression.py | 56 +++++++++++++++++-- tests/test_litellm/test_router.py | 7 ++- 3 files changed, 103 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index e7f58662249..b6f32c1c46d 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -135,19 +135,36 @@ def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: return None +def _compression_guardrail_classes() -> tuple[type, ...]: + """The registered guardrail classes whose provider compresses prompts.""" + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + return tuple(cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS) + + +def is_compression_guardrail(guardrail: object) -> bool: + """Whether `guardrail` is an instance of a compression guardrail provider. + + Both hops are validated through here. The two policy fields are operator-supplied + names and nothing else constrains them, so without this a name that resolves to an + ordinary guardrail would be handed the conversation and invoked: the routing hop + calls `apply_guardrail` directly, which POSTs the content wherever that guardrail + sends it, and the model hop is added to `metadata["guardrails"]`, which runs it even + when it is not `default_on`. + """ + classes: Final = _compression_guardrail_classes() + return bool(classes) and isinstance(guardrail, classes) + + def _active_compression_guardrails() -> tuple["CustomGuardrail", ...]: """Every currently-active guardrail whose type is a compression guardrail.""" import litellm from litellm.integrations.custom_guardrail import CustomGuardrail - from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry - compression_classes: Final = tuple( - cls for name, cls in guardrail_class_registry.items() if name in COMPRESSION_GUARDRAIL_PROVIDERS - ) - if not compression_classes: + if not _compression_guardrail_classes(): return () active: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) - return tuple(cb for cb in active if isinstance(cb, compression_classes) and cb.guardrail_name) + return tuple(cb for cb in active if is_compression_guardrail(cb) and cb.guardrail_name) async def arm_pre_call( @@ -192,7 +209,18 @@ async def arm_pre_call( ) ) - if policy.model is not None: + # Only a name that resolves to a real compression guardrail may be armed: this adds + # it to `metadata["guardrails"]`, which runs it even when it is not `default_on`. + armed_model_hop: Final = policy.model is not None and any( + guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails() + ) + if policy.model is not None and not armed_model_hop: + verbose_proxy_logger.warning( + "AutoRouter compression: '%s' is not an active compression guardrail; the model hop is uncompressed", + policy.model, + ) + + if armed_model_hop: _model_hop_armed.set(True) _, metadata = get_or_create_metadata_bucket(data) requested: Final = metadata.get("guardrails") @@ -249,6 +277,16 @@ async def messages_for_routing( ) return _as_routing_messages(messages) + # apply_guardrail below hands this guardrail the conversation and it POSTs the + # content to whatever service backs it, so the name has to be a compression + # guardrail rather than any guardrail the operator happened to name. + if not is_compression_guardrail(guardrail): + verbose_proxy_logger.warning( + "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages", + policy.routing, + ) + return _as_routing_messages(messages) + inputs: Final[GenericGuardrailAPIInputs] = { "structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape } diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 0b47e56cb02..676b3ba2967 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -186,15 +186,33 @@ class _RecordingCompressionGuardrail(CustomGuardrail): @pytest.fixture -def registered_guardrail(): +def registered_guardrail(monkeypatch): import litellm + from litellm.proxy.guardrails import guardrail_registry + # Registered under a compression provider name: both hops refuse a name that does + # not resolve to one, so a bare callback would (correctly) never be used. + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail) guardrail = _RecordingCompressionGuardrail(guardrail_name="fake-compress") litellm.logging_callback_manager.add_litellm_callback(guardrail) yield guardrail litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) +class _NonCompressionGuardrail(CustomGuardrail): + """A guardrail that is not a compression provider, e.g. a PII or content filter.""" + + def __init__(self, guardrail_name: str): + super().__init__(guardrail_name=guardrail_name) + self.called = False + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: str, logging_obj=None + ) -> GenericGuardrailAPIInputs: + self.called = True + return inputs + + class TestArmPreCall: @pytest.mark.asyncio async def test_no_router_is_noop(self): @@ -266,7 +284,13 @@ class TestArmPreCall: litellm.logging_callback_manager.remove_callback_from_all_lists(guardrail) @pytest.mark.asyncio - async def test_model_side_guardrail_is_requested_even_when_not_default_on(self): + async def test_model_side_guardrail_is_requested_even_when_not_default_on(self, monkeypatch): + import litellm + from litellm.proxy.guardrails import guardrail_registry + + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _RecordingCompressionGuardrail) + active = _RecordingCompressionGuardrail(guardrail_name="headroom-b") + litellm.logging_callback_manager.add_litellm_callback(active) router = _FakeRouter( [ { @@ -280,8 +304,11 @@ class TestArmPreCall: ] ) data = {"model": "smart-router", "messages": [{"role": "user", "content": "hi"}]} - await arm_pre_call(data=data, llm_router=router) - assert data["metadata"]["guardrails"] == ["headroom-b"] + try: + await arm_pre_call(data=data, llm_router=router) + assert data["metadata"]["guardrails"] == ["headroom-b"] + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(active) @pytest.mark.asyncio async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): @@ -347,6 +374,27 @@ class TestMessagesForRouting: assert result == [{"role": "user", "content": "[COMPRESSED] my ssn is [REDACTED]"}] assert registered_guardrail.request_data_seen[0]["messages"] == masked + @pytest.mark.asyncio + async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch): + """Regression (security): the policy fields are operator-supplied names that + nothing else constrains. apply_guardrail hands the guardrail the conversation + and it POSTs that content to whatever service backs it, so naming an ordinary + guardrail must not turn the routing hop into a way to ship prompts there.""" + import litellm + + other = _NonCompressionGuardrail(guardrail_name="pii-filter") + litellm.logging_callback_manager.add_litellm_callback(other) + try: + policy = AutoRouterCompressionPolicy(routing="pii-filter", model=None) + messages = [{"role": "user", "content": "my ssn is 123-45-6789"}] + + result = await messages_for_routing(policy=policy, messages=messages, request_kwargs={}) + + assert other.called is False + assert result == messages + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(other) + @pytest.mark.asyncio async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): """Regression: a real compression guardrail writes its stats onto whatever diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 9e6a88d3433..6279c8a5404 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -9686,7 +9686,12 @@ class TestAutoRouterCompressionDecoupling: return router, strategy @pytest.fixture - def registered_guardrail(self): + def registered_guardrail(self, monkeypatch): + from litellm.proxy.guardrails import guardrail_registry + + # Registered under a compression provider name: both hops refuse a name that + # does not resolve to one, so a bare callback would never be used. + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", self._CompressingGuardrail) guardrail = self._CompressingGuardrail(guardrail_name="fake-compress") litellm.logging_callback_manager.add_litellm_callback(guardrail) yield guardrail From 88985d00e2a1de43c616893141934cd0f444ce04 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 09:58:44 -0700 Subject: [PATCH 34/61] bump: litellm-enterprise 0.1.64 -> 0.1.65, litellm-proxy-extras 0.4.93 -> 0.4.94 --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 4 ++-- uv.lock | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index b6f482ccd86..3699087dbfa 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.64" +version = "0.1.65" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.64" +version = "0.1.65" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 97e9eb66bf2..82d31fec373 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.93" +version = "0.4.94" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.93" +version = "0.4.94" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index c1fde4af3f5..b889a3a0e60 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.93", - "litellm-enterprise==0.1.64", + "litellm-proxy-extras==0.4.94", + "litellm-enterprise==0.1.65", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/uv.lock b/uv.lock index 9e6343ff375..89205cd9527 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-01T21:00:02.682921Z" +exclude-newer = "2026-09-02T16:58:34.594994Z" exclude-newer-span = "P3D" [manifest] @@ -4771,12 +4771,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.64" +version = "0.1.65" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.93" +version = "0.4.94" source = { editable = "litellm-proxy-extras" } [[package]] From 1c16a5910b07415b2ab9cb6a54622f0296117f93 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 10:31:56 -0700 Subject: [PATCH 35/61] fix(tests): undo a stray whole-file reformat and arm a real guardrail test_router.py is not ruff-formatted on staging and CI's format check only scopes litellm/*.py, so running ruff format over the whole file rewrote ~900 lines of unrelated code. That reflow split long single-line patch() calls into multi-line form, which the test-quality gate counts individually, pushing TQ008 four over its ceiling. The file is back to staging's formatting with only the compression test class added. test_common_request_processing.py armed a model-side guardrail name with no such guardrail registered, which stopped working once both hops began requiring the name to resolve to an active compression guardrail. --- .../proxy/test_common_request_processing.py | 33 +- tests/test_litellm/test_router.py | 919 +++++++++++++----- 2 files changed, 675 insertions(+), 277 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index c0809e53d2e..96d9b0c5a26 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -383,6 +383,18 @@ class TestProxyBaseLLMRequestProcessing: """arm_pre_call must run before pre_call_hook: an auto router's own compression policy has to be in `data["metadata"]` (naming the model-side guardrail so it runs even if it isn't default_on) by the time guardrails see the request.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails import guardrail_registry + + # The model hop is only armed for a name that resolves to an active compression + # guardrail, so arming it has to have a real one to resolve to. + class _FakeCompressionGuardrail(CustomGuardrail): + pass + + monkeypatch.setitem(guardrail_registry.guardrail_class_registry, "headroom", _FakeCompressionGuardrail) + active_guardrail = _FakeCompressionGuardrail(guardrail_name="headroom-model") + litellm.logging_callback_manager.add_litellm_callback(active_guardrail) + processing_obj = ProxyBaseLLMRequestProcessing(data={}) mock_request = MagicMock(spec=Request) mock_request.headers = {} @@ -418,15 +430,18 @@ class TestProxyBaseLLMRequestProcessing: mock_proxy_config = MagicMock(spec=ProxyConfig) mock_proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None) - await processing_obj.common_processing_pre_call_logic( - request=mock_request, - general_settings={}, - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=mock_proxy_config, - route_type="acompletion", - llm_router=fake_llm_router, - ) + try: + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=fake_llm_router, + ) + finally: + litellm.logging_callback_manager.remove_callback_from_all_lists(active_guardrail) assert seen_metadata.get("guardrails") == ["headroom-model"] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6279c8a5404..d97fa7f2912 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -15,6 +15,7 @@ import pytest import respx + import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError @@ -136,18 +137,31 @@ def test_router_model_group_encrypted_content_affinity_callback_registration(): num_retries=0, ) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] - deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is False - assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config - assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < (litellm.callbacks.index(deployment_callback)) + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( + litellm.callbacks.index(deployment_callback) + ) router._add_encrypted_content_affinity_check(enable_global_affinity=True) callbacks = router.optional_callbacks or [] - encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] assert len(encrypted_content_callbacks) == 1 assert encrypted_content_callbacks[0].enable_global_affinity is True assert encrypted_content_callbacks[0].router is router @@ -178,9 +192,13 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + {model_group: ["encrypted_content_affinity"]} + ) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -221,7 +239,10 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] + assert ( + "encrypted_content_affinity_enabled" + not in disabled_request_kwargs["litellm_metadata"] + ) global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -241,7 +262,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert global_request_kwargs["litellm_metadata"][ + "encrypted_content_affinity_enabled" + ] @pytest.mark.asyncio @@ -288,10 +311,18 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) - encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) - assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) - assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callback) < ( + litellm.callbacks.index(deployment_callback) + ) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -302,7 +333,9 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -906,7 +939,9 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: + with patch.object( + litellm, "aretrieve_batch", return_value=AsyncMock() + ) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -927,7 +962,9 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: + with patch.object( + litellm, "afile_content", return_value=AsyncMock() + ) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -988,7 +1025,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: + with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1072,7 +1109,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert result is True, "Should return True when team_id and team_public_model_name match" + assert ( + result is True + ), "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1080,9 +1119,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert result is False, ( - "Should return False when team_id matches but model_name doesn't match team_public_model_name" - ) + assert ( + result is False + ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1098,18 +1137,30 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert result is True, "Should return True when team deployment has no team_public_model_name to match" + assert ( + result is True + ), "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) - assert result is True, "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id=None + ) + assert ( + result is True + ), "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") - assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id="any-team" + ) + assert ( + result is True + ), "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) + result = router.should_include_deployment( + model_name="different-model", model=deployment_without_team, team_id=None + ) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1118,7 +1169,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert result is True, "Should return True when matching model with exact model_name" + assert ( + result is True + ), "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1168,7 +1221,9 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + mock_response.text = ( + '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + ) with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1242,7 +1297,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: + with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1294,9 +1349,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: - with patch.object(router, "_get_client", return_value=None) as mock_get_client: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: + with patch.object( + router, "_get_client", return_value=None + ) as mock_get_client: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1331,7 +1392,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception, match="No deployment available") as exc_info: + with pytest.raises(Exception, match='No deployment available') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1359,9 +1420,15 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): mock_semaphore = asyncio.Semaphore(1) - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "_get_client", return_value=mock_semaphore) as mock_get_client: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "_get_client", return_value=mock_semaphore + ) as mock_get_client: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: result = await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_semaphore_function, @@ -1390,10 +1457,16 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): }, } - with patch.object(router, "_update_kwargs_with_deployment") as mock_update_kwargs: - with patch.object(router, "_get_client", return_value=None) as mock_get_client: - with patch.object(router, "async_routing_strategy_pre_call_checks") as mock_pre_call_checks: - with pytest.raises(Exception, match="Mock failure") as exc_info: + with patch.object( + router, "_update_kwargs_with_deployment" + ) as mock_update_kwargs: + with patch.object( + router, "_get_client", return_value=None + ) as mock_get_client: + with patch.object( + router, "async_routing_strategy_pre_call_checks" + ) as mock_pre_call_checks: + with pytest.raises(Exception, match='Mock failure') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -1461,9 +1534,9 @@ async def test_ageneric_api_call_deployment_model_overrides_alias(): original_generic_function=capture_model, ) - assert captured["model"] == "vertex_ai/gemini-2.5-flash", ( - f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" - ) + assert ( + captured["model"] == "vertex_ai/gemini-2.5-flash" + ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" @pytest.mark.asyncio @@ -1569,10 +1642,14 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) + access_groups = router.get_model_access_groups( + model_name="gpt-3.5-turbo", team_id=None + ) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") + access_groups = router.get_model_access_groups( + model_name="gpt-3.5-turbo", team_id="team_1" + ) assert list(access_groups.keys()) == ["default-models"] @@ -1667,7 +1744,9 @@ def test_model_group_info_cost_from_db_model_info(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._cached_get_model_group_info("my-custom-model") assert result is not None assert result.input_cost_per_token == 0.0001 @@ -1695,7 +1774,9 @@ def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._cached_get_model_group_info("my-custom-model-no-cost") assert result is not None assert result.input_cost_per_token is None @@ -1765,7 +1846,9 @@ def test_model_group_info_with_stringified_cost_values(): } return None - with patch.object(router, "get_deployment_model_info", side_effect=_model_info_with_str_costs): + with patch.object( + router, "get_deployment_model_info", side_effect=_model_info_with_str_costs + ): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -1811,7 +1894,9 @@ def test_model_group_info_db_fallback_with_stringified_cost_values(): ] ) - with patch.object(router, "get_deployment_model_info", side_effect=Exception("not found")): + with patch.object( + router, "get_deployment_model_info", side_effect=Exception("not found") + ): result = router._set_model_group_info( model_group="my-custom-model", user_facing_model_group_name="my-custom-model", @@ -2071,7 +2156,6 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] - async def _drain(): async for chunk in result: collected_chunks.append(chunk) @@ -2765,7 +2849,11 @@ def _make_responses_iterator( BaseResponsesAPIStreamingIterator, ) - base = LiteLLMCompletionStreamingIterator if bridge else BaseResponsesAPIStreamingIterator + base = ( + LiteLLMCompletionStreamingIterator + if bridge + else BaseResponsesAPIStreamingIterator + ) class _Iter(base): def __init__(self): @@ -2845,7 +2933,9 @@ async def test_aresponses_streaming_iterator_fallback(): BaseResponsesAPIStreamingIterator, ) - router = _make_router_with_fallback("anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6") + router = _make_router_with_fallback( + "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" + ) src = _make_responses_iterator( chunks=[MagicMock(type="response.created")], error=MidStreamFallbackError( @@ -2928,9 +3018,9 @@ async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback fbk = mock_fallback_utils.call_args.kwargs["kwargs"] assert "litellm_metadata" in fbk, "wrong metadata_variable_name" assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get("metadata", {}), ( - "model_group leaked into 'metadata' instead of 'litellm_metadata'" - ) + assert "model_group" not in fbk.get( + "metadata", {} + ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" @pytest.mark.asyncio @@ -3048,7 +3138,9 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): fallback_response_object = ResponsesAPIResponse( id="resp_test", created_at=0, model="gpt-4", object="response", output=[] ) - fallback_response_object.usage = ResponseAPIUsage(input_tokens=20, output_tokens=15, total_tokens=35) + fallback_response_object.usage = ResponseAPIUsage( + input_tokens=20, output_tokens=15, total_tokens=35 + ) fallback_event = ResponseCompletedEvent( type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=fallback_response_object, @@ -3057,7 +3149,9 @@ async def test_aresponses_streaming_iterator_combines_partial_usage(): with ( patch( "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace(usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4)), + return_value=SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) + ), ), patch.object( router, @@ -3358,7 +3452,9 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3386,10 +3482,14 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3416,10 +3516,14 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3446,7 +3550,9 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} + ) counting_threads = [] monkeypatch.setattr( @@ -3542,10 +3648,14 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3573,7 +3683,9 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) counting_threads = [] @@ -3608,7 +3720,9 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3644,7 +3758,9 @@ async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monke monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {}) calls = [] - monkeypatch.setattr(litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000) + monkeypatch.setattr( + litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 + ) count = await router._acount_pre_call_check_tokens( model="m", @@ -3672,7 +3788,9 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3697,7 +3815,9 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + ) deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, @@ -3739,7 +3859,9 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): ) assert with_instructions_tokens > input_only_tokens - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} + ) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( model="m", @@ -3808,7 +3930,9 @@ def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwarg prompt_only_tokens = router._count_pre_call_check_tokens( messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} + ) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -3928,7 +4052,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): + with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): router._count_pre_call_check_tokens(messages=None, input=None) @@ -3943,7 +4067,9 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): ], enable_pre_call_checks=True, ) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5}) + monkeypatch.setattr( + router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + ) counted: list[dict] = [] original = router._count_pre_call_check_tokens @@ -4033,7 +4159,9 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): + with patch.object( + litellm, "model_cost", {"test-custom-model": mock_custom_model_info} + ): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -4041,11 +4169,15 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model", model_name="test-model" + ) # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call + mock_get_model_info.assert_any_call( + model="gpt-3.5-turbo" + ) # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -4056,18 +4188,26 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) - assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) + assert ( + result["input_cost_per_token"] == 0.001 + ) # From custom model (overrides base 0.0015) + assert ( + result["output_cost_per_token"] == 0.002 + ) # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert result["mode"] == "chat" # From base model (overrides litellm "completion") + assert ( + result["mode"] == "chat" + ) # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) + assert ( + result["key"] == "gpt-3.5-turbo" + ) # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -4086,7 +4226,9 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model-no-base", model_name="test-model" + ) # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -4106,7 +4248,9 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") + result = router.get_deployment_model_info( + model_id="non-existent-model", model_name="test-model" + ) # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -4139,7 +4283,9 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") + result = router.get_deployment_model_info( + model_id="test-custom-model-invalid", model_name="test-model" + ) # Should handle exception gracefully and still return merged result assert result is not None @@ -4148,8 +4294,12 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): - result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") + with patch.object( + litellm, "get_model_info", side_effect=Exception("Not found") + ): + result = router.get_deployment_model_info( + model_id="non-existent", model_name="non-existent" + ) # Should return None when no model info is found assert result is None @@ -4172,7 +4322,9 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") + result = router.get_deployment_model_info( + model_id="custom-model-id", model_name="unknown-model" + ) # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -4208,11 +4360,15 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") + result = router.get_deployment_model_info( + model_id="custom-with-base", model_name="unknown-model" + ) # Should return custom_model_info merged with base model info assert result is not None - assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) + assert ( + result["input_cost_per_token"] == 0.01 + ) # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -4259,14 +4415,18 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): + with patch.object( + litellm, "model_cost", {"custom-model-id": mock_custom_model_info} + ): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") + result = router.get_deployment_model_info( + model_id="custom-model-id", model_name="test-model" + ) assert result is not None @@ -4276,17 +4436,29 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) - assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) + assert ( + result["input_cost_per_token"] == 0.01 + ) # From custom model (overrides base 0.03) + assert ( + result["max_tokens"] == 8000 + ) # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) - assert result["litellm_provider"] == "openai" # From base model (not in custom) - assert result["base_only_field"] == "base_value" # From base model (not in custom) + assert ( + result["output_cost_per_token"] == 0.06 + ) # From base model (not in custom) + assert ( + result["litellm_provider"] == "openai" + ) # From base model (not in custom) + assert ( + result["base_only_field"] == "base_value" + ) # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert result["mode"] == "completion" # From litellm model name info (not in custom or base) + assert ( + result["mode"] == "completion" + ) # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -4323,9 +4495,10 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", ( - f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] + == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" + ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" # Test Case 2: Bedrock invoke-with-response-stream endpoint kwargs = { @@ -4337,9 +4510,10 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="special-bedrock-model", model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", ) - assert result["endpoint"] == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream", ( - f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] + == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" + ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" # Test Case 3: Bedrock converse endpoint kwargs = { @@ -4351,9 +4525,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="bedrock-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse", ( - f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" + ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" # Test Case 4: Bedrock provider prefix auto-detected from model_name kwargs = { @@ -4364,9 +4538,9 @@ def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): model="router-model", model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", ) - assert result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke", ( - f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" - ) + assert ( + result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" + ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): @@ -4418,10 +4592,14 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [{"role": "user", "content": "This call should succeed by falling back."}] + messages = [ + {"role": "user", "content": "This call should succeed by falling back."} + ] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion(model="completely-unknown-model", messages=messages) + response = await router.acompletion( + model="completely-unknown-model", messages=messages + ) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -4513,10 +4691,15 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-claude-model" + ) assert credentials is not None - assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert ( + credentials["aws_bedrock_runtime_endpoint"] + == "https://bedrock-runtime.us-east-1.amazonaws.com" + ) assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -4543,7 +4726,9 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") + credentials = router.get_deployment_credentials_with_provider( + model_id="vertex-gemini" + ) assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -4583,7 +4768,9 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") + credentials = router.get_deployment_credentials_with_provider( + model_id="azure-gpt-4" + ) assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -4619,7 +4806,9 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-model" + ) assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -4662,7 +4851,9 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-model" + ) assert credentials is not None for key, value in aws_auth_params.items(): @@ -4697,11 +4888,15 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") + team_credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") + global_credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2" + ) assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -4742,11 +4937,15 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") + unscoped_credentials = router.get_deployment_credentials_with_provider( + model_id="gemini-2.5-pro" + ) assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") + owner_credentials = router.get_deployment_credentials_with_provider( + model_id="gemini-2.5-pro", team_id="team-b" + ) assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -4773,8 +4972,16 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None - assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None + assert ( + router.get_deployment_credentials_with_provider( + model_id="gemini-2.5-pro", team_id="team-a" + ) + is None + ) + assert ( + router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") + is None + ) def test_deployment_usable_by_team_helpers(): @@ -4814,7 +5021,9 @@ def test_deployment_usable_by_team_helpers(): assert router._deployment_usable_by_team(shared, "team-a") is True assert router._deployment_usable_by_team(shared, None) is True - picked = router._get_model_group_deployment_usable_by_team(model_group_name="gemini-2.5-pro", team_id="team-a") + picked = router._get_model_group_deployment_usable_by_team( + model_group_name="gemini-2.5-pro", team_id="team-a" + ) assert picked is not None assert picked.litellm_params.vertex_project == "shared-project" @@ -4824,7 +5033,12 @@ def test_deployment_usable_by_team_helpers(): assert owner_picked is not None assert owner_picked.litellm_params.vertex_project == "team-b-project" - assert router._get_model_group_deployment_usable_by_team(model_group_name="unknown-model", team_id="team-a") is None + assert ( + router._get_model_group_deployment_usable_by_team( + model_group_name="unknown-model", team_id="team-a" + ) + is None + ) def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -4856,7 +5070,9 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") + owner_credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-b" + ) assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -4868,11 +5084,21 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) + is not None + ) router.delete_deployment(id="team-wildcard-id") - assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) + is None + ) def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -4945,13 +5171,22 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) - credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") + router.upsert_deployment( + deployment=Deployment(**_team_wildcard_model(api_key="new-key")) + ) + credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-1" + ) + is None + ) def test_get_available_guardrail_single_deployment(): @@ -5140,7 +5375,9 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ), model="gpt-3.5-turbo", model_id="model-123", model_group="openai-gpt", @@ -5219,8 +5456,12 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" - assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" + assert ( + cached_result is not None + ), "Model ID should be cached for anthropic_messages call type" + assert ( + cached_result["model_id"] == test_model_id + ), f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -5245,7 +5486,9 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Deployment tags should be propagated to kwargs metadata @@ -5274,7 +5517,9 @@ def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): # Simulate request that already has tags (from request body or key/team level) kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Both sources should be merged, no duplicates @@ -5301,7 +5546,9 @@ def test_update_kwargs_with_deployment_no_tags(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + deployment = router.get_deployment_by_model_group_name( + model_group_name="gpt-4o-mini" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # No tags key should be added if deployment has no tags @@ -5339,7 +5586,9 @@ def test_update_kwargs_with_deployment_merges_tools(): }, ], } - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Tools should be merged: deployment first, then request @@ -5370,7 +5619,9 @@ def test_update_kwargs_with_deployment_merge_tools_deployment_only(): ) kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["tools"] == [{"type": "web_search"}] @@ -5399,7 +5650,9 @@ def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice "metadata": {}, "tool_choice": "none", } - deployment = router.get_deployment_by_model_group_name(model_group_name="o3-deep-research") + deployment = router.get_deployment_by_model_group_name( + model_group_name="o3-deep-research" + ) router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # Request tool_choice should be preserved (merged tools still applied) @@ -5501,8 +5754,12 @@ def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name="generic_api_call" + ) assert "litellm_metadata" in kwargs model_info = kwargs["litellm_metadata"]["model_info"] @@ -5534,8 +5791,12 @@ def test_update_kwargs_with_deployment_model_info_in_metadata(): ) kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name(model_group_name="claude-sonnet-4") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=None) + deployment = router.get_deployment_by_model_group_name( + model_group_name="claude-sonnet-4" + ) + router._update_kwargs_with_deployment( + deployment=deployment, kwargs=kwargs, function_name=None + ) assert "metadata" in kwargs model_info = kwargs["metadata"]["model_info"] @@ -5648,7 +5909,6 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - async def _drain(): async for chunk in result: collected.append(chunk) @@ -5675,7 +5935,6 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - async def _drain(): async for chunk in result: collected.append(chunk) @@ -5892,17 +6151,23 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" + assert ( + len(deployment_ids) == 2 + ), "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="other-team" + ) assert len(deployments) == 0 @@ -5947,8 +6212,12 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") - assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert ( + len(deployments) == 2 + ), "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -6073,7 +6342,9 @@ def test_explicit_model_access_does_not_force_access_group_filtering(): }, ) - deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] + deployment_groups = [ + d.get("model_info", {}).get("access_groups") for d in deployments + ] assert ["AG1"] in deployment_groups assert ["AG2"] in deployment_groups @@ -6118,7 +6389,9 @@ def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6193,7 +6466,9 @@ def test_access_group_block_does_not_silently_use_default_fallback_model( orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6260,7 +6535,9 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba orig_groups = router.get_model_access_groups - def fake_get_model_access_groups(model_name=None, model_access_group=None, team_id=None): + def fake_get_model_access_groups( + model_name=None, model_access_group=None, team_id=None + ): if model_name == "gpt-5" and model_access_group is None: return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} return orig_groups( @@ -6315,7 +6592,9 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ) assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names(model="gpt-5", request_team_id=None) + router_in_names._try_early_resolve_deployments_for_model_not_in_names( + model="gpt-5", request_team_id=None + ) is None ) assert ( @@ -6337,8 +6616,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): ] ) - pattern_result = pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None + pattern_result = ( + pattern_router._try_early_resolve_deployments_for_model_not_in_names( + model="openai/gpt-4o-mini", request_team_id=None + ) ) assert pattern_result is not None resolved_model, pattern_deployments = pattern_result @@ -6364,8 +6645,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): }, } - default_result = default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None + default_result = ( + default_router._try_early_resolve_deployments_for_model_not_in_names( + model="brand-new-model", request_team_id=None + ) ) assert default_result is not None resolved_model, default_deployment = default_result @@ -6373,7 +6656,10 @@ def test_try_early_resolve_deployments_for_model_not_in_names(): assert isinstance(default_deployment, dict) assert default_deployment["litellm_params"]["model"] == "brand-new-model" # The original default_deployment must not be mutated. - assert default_router.default_deployment["litellm_params"]["model"] == "openai/will-be-overridden" + assert ( + default_router.default_deployment["litellm_params"]["model"] + == "openai/will-be-overridden" + ) def _router_with_two_deployments(blocked_flags): @@ -6421,7 +6707,10 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} + { + uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} + for uid in unhealthy_ids + } ) @@ -6492,7 +6781,9 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol @pytest.mark.asyncio async def test_async_get_healthy_deployments_skips_blocked_deployment(): router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments(model="gpt-4o", parent_otel_span=None) + healthy, all_dep = await router._async_get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" not in healthy_ids assert "dep-1" in healthy_ids @@ -6501,7 +6792,9 @@ async def test_async_get_healthy_deployments_skips_blocked_deployment(): def test_get_healthy_deployments_sync_skips_blocked_deployment(): router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments(model="gpt-4o", parent_otel_span=None) + healthy, all_dep = router._get_healthy_deployments( + model="gpt-4o", parent_otel_span=None + ) healthy_ids = [d["model_info"]["id"] for d in healthy] assert "dep-0" in healthy_ids assert "dep-1" not in healthy_ids @@ -6518,7 +6811,9 @@ def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) + deployments = await router.async_get_healthy_deployments( + model="gpt-4o", request_kwargs={} + ) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -6560,7 +6855,9 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) + deployment = router.get_available_deployment_for_pass_through( + model="gpt-4o", request_kwargs={} + ) assert deployment["model_info"]["id"] == "pt-1" @@ -6569,7 +6866,9 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) + router.get_available_deployment_for_pass_through( + model="pt-0", request_kwargs={} + ) def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): @@ -6593,7 +6892,9 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] + assert [m["model_info"]["id"] for m in router.get_model_list()] == [ + "bedrock-iam-pt" + ] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -6604,7 +6905,12 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" + assert ( + passthrough_router.get_credentials( + custom_llm_provider="openai", region_name=None + ) + == "sk-fake-for-tests" + ) def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -6638,9 +6944,16 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() - assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False assert ( - litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=missing_blocked) + ) + is False + ) + assert ( + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) is True ) @@ -6680,7 +6993,9 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): + def test_request_timeout_stored_independently_when_both_set( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -6698,16 +7013,22 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): + def test_non_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - def test_stream_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): + def test_stream_prefers_request_timeout_over_router_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) # stream=True resolves through _get_stream_timeout; request_timeout must win. assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - def test_explicit_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): + def test_explicit_stream_timeout_still_wins_over_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330, stream_timeout=45) assert router._get_stream_timeout(kwargs={}, data={}) == 45 @@ -6723,13 +7044,22 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_per_deployment_timeout_overrides_request_timeout(self, explicit_request_timeout): + def test_per_deployment_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - def test_per_request_timeout_overrides_request_timeout(self, explicit_request_timeout): + def test_per_request_timeout_overrides_request_timeout( + self, explicit_request_timeout + ): router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={"timeout": 60}, data={"timeout": 120}) == 60 + assert ( + router._get_non_stream_timeout( + kwargs={"timeout": 60}, data={"timeout": 120} + ) + == 60 + ) # --------------------------------------------------------------------------- @@ -7038,7 +7368,9 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) + active = router.cooldown_cache.get_active_cooldowns( + model_ids=["dep-1"], parent_otel_span=None + ) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -7047,7 +7379,12 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True + assert ( + router.deployment_callback_on_failure( + self._kwargs(self._auth_error()), None, now, now + ) + is True + ) assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -7062,7 +7399,12 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False + assert ( + router.deployment_callback_on_failure( + self._kwargs(exception), None, now, now + ) + is False + ) assert "dep-1" not in self._cooled_down_ids(router) @@ -7119,13 +7461,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta( - images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] - ) + images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) + annotations_delta = Delta( + annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] + ) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -7169,8 +7511,12 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): - assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert router.get_configured_token_limits( + "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + ) == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -7243,8 +7589,13 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): - assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -7303,8 +7654,13 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): - assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) def test_get_configured_display_name_treats_malformed_values_as_absent(): @@ -7537,16 +7893,13 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with ( - patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, - patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, - ), + with patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, ): await router.acreate_batch( model="bedrock-batch-model", @@ -7575,13 +7928,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + expected_response = VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch( - "litellm.vector_stores.main.asearch", new=mock_asearch - ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7615,9 +7968,7 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch( - "litellm.vector_stores.main.acreate", new=mock_acreate - ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -7633,13 +7984,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + expected_response = VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch( - "litellm.vector_stores.main.search", new=mock_search - ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -7648,7 +7999,9 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") + search_response = router.vector_store_search( + vector_store_id="v", query="q", custom_llm_provider="s3_vectors" + ) assert search_response is expected_response mock_search.assert_called_once() @@ -7662,9 +8015,7 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch( - "litellm.vector_stores.main.create", new=mock_create - ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -7699,7 +8050,9 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + }, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -8004,7 +8357,9 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) + **self._hybrid_router_params( + {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + ) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -8113,7 +8468,9 @@ class TestPreRoutingStrategyRegistryLifecycle: ({"model": "openai/gpt-4o"}, False), ] for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing(litellm_params=LiteLLM_Params(**params)) + actual = router._deployment_participates_in_adaptive_routing( + litellm_params=LiteLLM_Params(**params) + ) assert actual is expected, params["model"] @@ -8300,16 +8657,22 @@ class TestUpsertDeploymentRollback: router.delete_deployment(id="prod-1") assert router.has_model_id("prod-1") is False - router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") + router._restore_deployment_after_failed_upsert( + previous_deployment=previous, model_id="prod-1" + ) restored = router.get_deployment(model_id="prod-1") assert restored is not None assert restored.litellm_params.model == "gpt-4o" - router._restore_deployment_after_failed_upsert(previous_deployment=previous, model_id="prod-1") + router._restore_deployment_after_failed_upsert( + previous_deployment=previous, model_id="prod-1" + ) assert len(router.model_list) == 1 - router._restore_deployment_after_failed_upsert(previous_deployment=None, model_id="prod-1") + router._restore_deployment_after_failed_upsert( + previous_deployment=None, model_id="prod-1" + ) assert len(router.model_list) == 1 @@ -9075,14 +9438,18 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], + "candidates": [ + {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} + ], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) + @pytest.mark.parametrize( + "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] + ) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -9205,7 +9572,9 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) + await router.acompletion( + model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs + ) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -9244,9 +9613,9 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop] assert breadcrumbs, "no retry breadcrumbs were recorded" - assert any("fallback_depth" in breadcrumb for breadcrumb in breadcrumbs), ( - "no breadcrumb carried router walk state, so this test cannot see the leak" - ) + assert any( + "fallback_depth" in breadcrumb for breadcrumb in breadcrumbs + ), "no breadcrumb carried router walk state, so this test cannot see the leak" for breadcrumb in breadcrumbs: assert "attempted_targets" not in breadcrumb @@ -9292,9 +9661,7 @@ async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_ke breadcrumbs = metadata["previous_models"] assert breadcrumbs, "no retry breadcrumbs were recorded" dumped = json.dumps(breadcrumbs, default=str) - assert container_key in dumped, ( - "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" - ) + assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped @@ -9399,7 +9766,9 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), + mock_response=litellm.InternalServerError( + message=huge_message, llm_provider="openai", model="group-a" + ), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -9465,7 +9834,9 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): try: router._ensure_deployment_affinity_callback() router._ensure_deployment_affinity_callback() - affinity_callbacks = [cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)] + affinity_callbacks = [ + cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) + ] assert len(affinity_callbacks) == 1 finally: for cb in router.optional_callbacks or []: @@ -9607,7 +9978,9 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) + response = await router.acompletion( + model="smart-alias", messages=self._messages(), metadata=metadata + ) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -9829,6 +10202,8 @@ class TestAutoRouterCompressionDecoupling: assert registered_guardrail.call_count == 0 +@pytest.mark.usefixtures("local_model_cost_map") + @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -9855,14 +10230,17 @@ class TestAzureBaseModelFallbackLogging: def test_map_known_deployment_name_resolves_without_error_log(self): router = self._router_with_azure_deployment("azure/gpt-4o") - with patch("litellm.router.verbose_router_logger.error") as mock_error: + with patch( + "litellm.router.verbose_router_logger.error" + ) as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert not any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( - f"unexpected error log: {mock_error.call_args_list}" - ) + assert not any( + "Could not identify azure model" in str(call) + for call in mock_error.call_args_list + ), f"unexpected error log: {mock_error.call_args_list}" # the fallback resolution must actually surface the map values assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] @@ -9870,14 +10248,17 @@ class TestAzureBaseModelFallbackLogging: def test_unmappable_deployment_name_still_logs_error(self): router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - with patch("litellm.router.verbose_router_logger.error") as mock_error: + with patch( + "litellm.router.verbose_router_logger.error" + ) as mock_error: model_info = router.get_router_model_info( deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) - assert any("Could not identify azure model" in str(call) for call in mock_error.call_args_list), ( - "expected the error log for an unmappable azure deployment name" - ) + assert any( + "Could not identify azure model" in str(call) + for call in mock_error.call_args_list + ), "expected the error log for an unmappable azure deployment name" # unmappable names resolve to a zeroed stub — unchanged behavior assert model_info.get("max_input_tokens") is None @@ -9904,7 +10285,6 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] - def test_model_group_info_intersects_supported_reasoning_efforts(): router = litellm.Router( model_list=[ @@ -9995,6 +10375,7 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o assert result.supported_reasoning_efforts is None + def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose registry entry declares parallel function calling must flip the group to True instead of False.""" @@ -10227,7 +10608,6 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" - # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -10338,7 +10718,10 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' + ) def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -10383,7 +10766,9 @@ async def test_anthropic_messages_streaming_iterator_passthrough(): [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] ) - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] @@ -10402,14 +10787,12 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] ) - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) collected = [chunk async for chunk in wrapped] - assert collected == [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_content_chunk("hi"), - message_stop, - ] + assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] @pytest.mark.asyncio @@ -10421,7 +10804,9 @@ async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) collected = [chunk async for chunk in wrapped] assert collected == [_anthropic_messages_message_start_chunk(), message_stop] @@ -10492,9 +10877,7 @@ async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): yield _anthropic_messages_message_start_chunk() yield _anthropic_messages_content_chunk("hi") - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source(), initial_kwargs={"model": "primary"} - ) + wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() content_released.set() @@ -12295,9 +12678,9 @@ class TestTierParamsTheTargetAccepts: def test_declared_param_allowlist_ignores_malformed_declarations(self): """A str is iterable, so without the type guard a YAML scalar mistake like allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist( - {"allowed_openai_params": ["reasoning_effort", 3]} - ) == frozenset({"reasoning_effort"}) + assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( + {"reasoning_effort"} + ) assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() assert litellm.Router._declared_param_allowlist({}) == frozenset() @@ -12377,11 +12760,7 @@ class TestTierParamsTheTargetAccepts: @pytest.mark.parametrize( "deployment", - [ - {"model_name": "x"}, - {"model_name": "x", "litellm_params": {}}, - {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}, - ], + [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], ) def test_deployment_accepts_param_fails_open(self, deployment): """An unresolvable deployment must not be the reason a param is dropped.""" @@ -12610,7 +12989,9 @@ class TestPreRoutingTierDrivesFallbacks: async def test_the_selected_tier_fallback_chain_runs(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + response = await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}] + ) assert response.choices[0].message.content == "from backup-a" @@ -12643,7 +13024,9 @@ class TestPreRoutingTierDrivesFallbacks: async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): router = self._router([{"tier1": ["backup-a"]}]) - response = await router.acompletion(model="tier1", messages=[{"role": "user", "content": "hi"}]) + response = await router.acompletion( + model="tier1", messages=[{"role": "user", "content": "hi"}] + ) assert response.choices[0].message.content == "from backup-a" From 98784360e85186f798c7ffac797aba4020c964fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 10:37:14 -0700 Subject: [PATCH 36/61] test(e2e): cover Anthropic /chat/completions streaming and tool calls Adds TestAnthropicChatCompletions to the chat completions regression suite, registering a claude-haiku-4-5 deployment via /model/new and asserting the streamed call delivers real content deltas and a tool-forced call returns a well-formed get_weather tool_call on both the non-streamed and streamed paths. Covers three P0 registry cells that had no e2e test. --- .../test_chat_completions_regression_e2e.py | 98 ++++++++++++++++++- 1 file changed, 97 insertions(+), 1 deletion(-) diff --git a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py index 68c0dfab897..87bd32d8dab 100644 --- a/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_regression_e2e.py @@ -11,7 +11,7 @@ fails that provider's row here. The per-provider classes below cover the OpenAI-compatible /chat/completions translation for providers customers reach by registering their own deployment -via /model/new (Cohere, Gemini, hosted_vllm), each deleted on teardown. +via /model/new (Cohere, Gemini, hosted_vllm, Anthropic), each deleted on teardown. """ from __future__ import annotations @@ -46,6 +46,7 @@ pytestmark = pytest.mark.e2e COHERE_BACKEND = "cohere/command-r-08-2024" GEMINI_BACKEND = "gemini/gemini-2.5-flash" OPENAI_BACKEND = "openai/gpt-5.6" +ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5-20251001" BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" @@ -746,3 +747,98 @@ class TestBedrockConverseChatCompletions: response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32))) _assert_describes_cat(response) + + +class TestAnthropicChatCompletions: + """Anthropic via the OpenAI-compatible /chat/completions path, the translation + customers on the OpenAI SDK rely on when they route to Claude. The streamed call + must deliver real content deltas, and a tool-forced call must come back as a + well-formed tool_call on both the non-streamed and streamed paths. + """ + + def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str: + model = f"{prefix}-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + @pytest.mark.covers( + "llm.chat_completions.anthropic.basic.stream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_streams_real_content( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], + max_tokens=64, + stream=True, + ), + ) + _assert_streamed_completion(result) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.tool_use.nonstream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_returns_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-tool") + key = resources.key() + + response = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + ), + ) + ) + _assert_weather_tool_call(response) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.tool_use.stream.works", + exercised_on=["chat_completions"], + ) + def test_anthropic_chat_streams_tool_call( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = self._register(client, resources, "e2e-anthropic-tool-stream") + key = resources.key() + + result = client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.") + ], + tools=[_WEATHER_TOOL], + tool_choice="required", + max_tokens=128, + stream=True, + ), + ) + assert result.ok and result.is_streaming, f"tool stream was not established: {result}" + assert result.stream_error is None, f"tool stream carried an error event: {result.stream_error}" + name, arguments = _streamed_tool_call(result.stream_events) + assert name == "get_weather", f"streamed tool call named {name!r}: {result.stream_events[:5]}" + args = _WeatherArgs.model_validate_json(arguments) + assert args.location.strip(), f"streamed tool call arguments missing location: {arguments!r}" From df544fcc532179e3ff388a41032a514ce41c5020 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 10:46:05 -0700 Subject: [PATCH 37/61] test(e2e): cover key spend reset, regenerate grace period, and the llm_api_routes grant Three deterministic proxy-only cells from the coverage registry that had no e2e test. A key over its max_budget is reset to 0 through /key/{key}/reset_spend and must both read back 0 on /key/info and serve traffic again. /key/regenerate with grace_period keeps the old key valid until the period elapses and rejects it 401 afterwards. A key whose allowed_routes is the llm_api_routes group must reach /chat/completions and /embeddings while /model/new stays 403. KeyRegenerateBody gains grace_period and the management client gains reset_key_spend so the tests stay on the shared typed transport. --- .../access_control/test_access_control_e2e.py | 28 ++++++++- tests/e2e/management/management_client.py | 16 ++++- .../e2e/management/test_key_management_e2e.py | 59 ++++++++++++++++++- tests/e2e/management/test_management_e2e.py | 34 +++++++++++ tests/e2e/models.py | 10 ++++ 5 files changed, 143 insertions(+), 4 deletions(-) diff --git a/tests/e2e/access_control/test_access_control_e2e.py b/tests/e2e/access_control/test_access_control_e2e.py index af7e9a099fd..c30dadc49ae 100644 --- a/tests/e2e/access_control/test_access_control_e2e.py +++ b/tests/e2e/access_control/test_access_control_e2e.py @@ -24,7 +24,7 @@ from access_control_client import ( from e2e_config import unique_marker from e2e_http import Success, UnauthorizedError, UnknownApiError, unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatResponse, EmbedBody, LiteLLMParamsBody from proxy_client import ProxyClient pytestmark = pytest.mark.e2e @@ -32,6 +32,7 @@ pytestmark = pytest.mark.e2e ALLOWED_MODEL = "gemini-2.5-flash" DISALLOWED_MODEL = "gpt-5.5" VIRTUAL_KEY_BACKEND = "anthropic/claude-haiku-4-5-20251001" +EMBEDDING_MODEL = "openai-text-embedding-3-small" class TestAccessControl: @@ -71,6 +72,31 @@ class TestAccessControl: f"403 body must be a model-access denial, got: {result.body[:300]}" ) + @pytest.mark.covers("other.auth.virtual_key.route_group_allowed") + def test_llm_api_routes_group_grants_every_llm_endpoint( + self, client: AccessControlClient, resources: ResourceManager + ) -> None: + """allowed_routes=["llm_api_routes"] names a route group, not a path: one + entry must open every LLM endpoint while the management routes stay shut.""" + key = client.llm_only_key() + resources.defer(lambda: client.delete_key(key)) + + chat = client.chat_status(key, ALLOWED_MODEL, f"capital of France? {unique_marker()}") + assert chat.status_code == 200, ( + f"llm_api_routes key must reach /chat/completions, got {chat.status_code}: {chat.body[:300]}" + ) + assert ChatResponse.model_validate_json(chat.body).choices, ( + f"200 must carry a real completion, not an error envelope: {chat.body[:300]}" + ) + + embedding = unwrap(client.proxy.embed(key, EmbedBody(model=EMBEDDING_MODEL, input=f"route group {unique_marker()}"))) + assert embedding.model, f"llm_api_routes key reached /embeddings but got no model back: {embedding}" + + denied = client.create_model_status(key, f"e2e-route-group-{unique_marker()}") + assert denied.status_code == 403 and ROUTE_NOT_ALLOWED_MARKER in denied.body, ( + f"the same key must still be shut out of /model/new, got {denied.status_code}: {denied.body[:300]}" + ) + def test_llm_only_key_forbidden_from_management_route_403( self, client: AccessControlClient, resources: ResourceManager ) -> None: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 387280c8023..2b897f5f07f 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -40,6 +40,8 @@ from models import ( KeyListParams, KeyListResponse, KeyRegenerateBody, + KeyResetSpendBody, + KeyResetSpendResponse, KeyUpdateBody, ModelDeleteBody, OrgDeleteBody, @@ -191,16 +193,26 @@ class ManagementClient: response_type=NoBody, ) ) - def regenerate_key(self, key: str) -> str: + def regenerate_key(self, key: str, *, grace_period: str | None = None) -> str: return unwrap( self.proxy.transport.post( "/key/regenerate", headers=self.proxy.transport.master, - json=KeyRegenerateBody(key=key), + json=KeyRegenerateBody(key=key, grace_period=grace_period), response_type=KeyGenerateResponse, ) ).key + def reset_key_spend(self, key: str, reset_to: float) -> KeyResetSpendResponse: + return unwrap( + self.proxy.transport.post( + f"/key/{key}/reset_spend", + headers=self.proxy.transport.master, + json=KeyResetSpendBody(reset_to=reset_to), + response_type=KeyResetSpendResponse, + ) + ) + def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]: """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is who is asking: the master key by default, or a virtual key.""" diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index 711175abb0d..d4347b8c0e7 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -18,7 +18,7 @@ from typing import Literal import pytest from e2e_config import unique_marker -from e2e_http import NoBody, unwrap +from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient from models import KeyDeleteBody, KeyGenerateBody, KeyUpdateBody @@ -26,6 +26,9 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e +TINY_BUDGET = 3e-6 +SPEND_MODEL = "claude-haiku-4-5" + class KeyToggleBlockBody(BaseModel): key: str @@ -82,6 +85,34 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke return key +def _is_budget_block(outcome: StreamingResponse) -> bool: + return not outcome.ok and "budget_exceeded" in outcome.body + + +def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None: + """Drive paid calls until the key's max_budget refuses one. The first call spends, + the reservation counter trips the cap, and the next call is the 429.""" + for _ in range(40): + outcome = client.chat_status(key, SPEND_MODEL, f"spend {unique_marker()}") + if _is_budget_block(outcome): + assert outcome.status_code == 429, ( + f"budget refusal must be 429, got {outcome.status_code}: {outcome.body[:200]}" + ) + return + assert outcome.ok, f"paid call failed before the budget tripped ({outcome.status_code}): {outcome.body[:300]}" + time.sleep(2) + pytest.fail(f"max_budget={TINY_BUDGET} never blocked a call on the key") + + +def _settled_spend(client: ManagementClient, key: str) -> float | None: + """The key's recorded spend once it is positive and unchanged across two reads a + poll interval apart, so no batched spend write is still in flight when we reset.""" + first = client.proxy.key_info(key).spend or 0.0 + time.sleep(client.proxy.poll_interval) + second = client.proxy.key_info(key).spend or 0.0 + return second if first > 0 and first == second else None + + def _block(client: ManagementClient, key: str) -> None: _ = unwrap( client.proxy.transport.post( @@ -197,6 +228,32 @@ class TestKeyManagementRoutes: "/key/info never reported max_budget 42.0 after /key/bulk_update before the deadline", ) + @pytest.mark.covers("other.key_mgmt.spend_reset.resets_to_value") + def test_reset_spend_zeroes_recorded_spend_and_lifts_the_budget_block( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + key = _generate_key(client, resources, KeyGenerateBody(models=[SPEND_MODEL], max_budget=TINY_BUDGET)) + _spend_until_budget_blocks(client, key) + recorded = _poll( + client, lambda: _settled_spend(client, key), "key spend never landed in /key/info before the deadline" + ) + + reset = client.reset_key_spend(key, reset_to=0.0) + assert reset.previous_spend == recorded, ( + f"reset_spend reported previous_spend {reset.previous_spend}, /key/info had recorded {recorded}" + ) + assert reset.spend == 0.0, f"reset_spend to 0 reported spend {reset.spend}" + assert client.proxy.key_info(key).spend == 0.0, "/key/info still reports spend after the reset to 0" + + def call_allowed_again() -> bool | None: + outcome = client.chat_status(key, SPEND_MODEL, f"after reset {unique_marker()}") + if _is_budget_block(outcome): + return None + assert outcome.ok, f"post-reset call failed ({outcome.status_code}): {outcome.body[:300]}" + return True + + _ = _poll(client, call_allowed_again, "the key stayed budget-blocked after its spend was reset to 0") + @pytest.mark.covers("mgmt.key.generate.admin_only") def test_generate_forbidden_for_non_admin_key( self, client: ManagementClient, resources: ResourceManager diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index a56eb853823..eace8f2e3b4 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -42,6 +42,10 @@ from models import ( pytestmark = pytest.mark.e2e +REGENERATE_GRACE_PERIOD = "15s" +REGENERATE_GRACE_SECONDS = 15.0 + + def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: @@ -365,6 +369,36 @@ class TestKeyRegeneration: client, old_rejected, "old key was still accepted after regeneration (never rejected 401) at the deadline" ) + @pytest.mark.covers("other.key_mgmt.regenerate.grace_period_honored") + def test_regenerate_with_grace_period_keeps_old_key_until_revoked( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + old_key = _generate_key(client, resources, KeyGenerateBody(models=["gpt-5.5"])) + + new_key = client.regenerate_key(old_key, grace_period=REGENERATE_GRACE_PERIOD) + resources.defer(lambda: client.proxy.delete_key(new_key)) + revoke_at = time.monotonic() + REGENERATE_GRACE_SECONDS + assert new_key != old_key, "regenerate returned the same key string, so no rotation happened" + + def old_accepted() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code != 401 else None + + _ = _poll(client, old_accepted, "old key was rejected 401 inside its grace period at the deadline") + assert time.monotonic() < revoke_at, ( + f"old key was only accepted after its {REGENERATE_GRACE_PERIOD} grace period had elapsed" + ) + + def old_rejected() -> bool | None: + outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") + return True if outcome.status_code == 401 else None + + _ = _poll( + client, + old_rejected, + f"old key was still accepted past its {REGENERATE_GRACE_PERIOD} grace period (never 401) at the deadline", + ) + class TestTeamRoutes: @pytest.mark.covers("mgmt.team.new.persists") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index b5229744d6f..5de49ead3ed 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -83,6 +83,16 @@ class KeyGenerateResponse(BaseModel): class KeyRegenerateBody(BaseModel): key: str + grace_period: str | None = None + + +class KeyResetSpendBody(BaseModel): + reset_to: float + + +class KeyResetSpendResponse(BaseModel): + spend: float + previous_spend: float class KeyDeleteBody(BaseModel): From 88c46fb1defa968348b9e780bd63453f5f8f3e94 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 10:52:46 -0700 Subject: [PATCH 38/61] test(e2e): cover Anthropic and OpenAI prompt caching, Cohere embeddings, and costed /openai chat passthrough Four registry cells that had no e2e test. The cache_control suite gains a direct Anthropic case (the same cache_control prefix the Bedrock and Vertex rows send) and an OpenAI case, where caching is automatic so the prefix goes out as a plain system string with a prompt_cache_key; both assert the second identical call reports cache-read tokens. The shared second-call helper now takes the send callable so the OpenAI shape fits without a second copy of the retry loop. The embeddings suite gains a cohere/embed-v4.0 deployment that must return a non-zero vector, and the passthrough suite gains an OpenAI-format chat through the raw /openai/v1/chat/completions prefix that must relay a real completion and log a costed pass_through_endpoint row whose token counts match the usage the caller was served. --- .../e2e/llm_translation/test_cache_control.py | 76 +++++++++++++++++-- .../test_embeddings_endpoint_e2e.py | 22 +++++- .../llm_translation/test_passthrough_e2e.py | 34 ++++++++- 3 files changed, 124 insertions(+), 8 deletions(-) diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index a2c17b0fb66..0d224061381 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -10,6 +10,11 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix intentionally not covered here. - Vertex (gemini-2.5-flash): prompt caching via ``cache_control`` context caching; the second identical call must report cached prompt tokens > 0. +- Anthropic (claude-haiku-4-5, direct): the same ``cache_control`` prefix over + the OpenAI-compatible route; the second call must report cache-read tokens > 0. +- OpenAI (gpt-5.6): automatic prompt caching needs no request marker, so the + cacheable prefix goes out as a plain system string with a ``prompt_cache_key`` + and the second call must report ``prompt_tokens_details.cached_tokens`` > 0. service_tier lives in test_provider_features_e2e.py. @@ -21,6 +26,7 @@ built from the typed content blocks shared in ``endpoints_client.py``. from __future__ import annotations import time +from collections.abc import Callable import pytest from pydantic import BaseModel @@ -29,7 +35,7 @@ from e2e_config import unique_marker from e2e_http import Result, unwrap from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager -from models import ChatResponse, LiteLLMParamsBody, Usage +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage from passthrough_client import PassthroughClient import os @@ -37,6 +43,8 @@ pytestmark = pytest.mark.e2e BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" VERTEX_MODEL = "vertex_ai/gemini-2.5-flash" +ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001" +OPENAI_MODEL = "openai/gpt-5.6" class CacheChatBody(BaseModel): @@ -89,17 +97,36 @@ def _cache_chat( ) +def _plain_cache_chat( + client: PassthroughClient, key: str, model: str, prefix: str, cache_key: str +) -> Result[ChatResponse]: + """The same cacheable prefix as a plain system string, for providers that cache + automatically and take no per-block marker (OpenAI).""" + return client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ + ChatMessage(role="system", content=prefix), + ChatMessage(role="user", content="Reply with one word."), + ], + max_tokens=64, + prompt_cache_key=cache_key, + ), + ) + + def _assert_cache_read_on_second_call( - client: PassthroughClient, key: str, model: str + model: str, send: Callable[[str], Result[ChatResponse]] ) -> None: prefix = _cacheable_prefix() - first = unwrap(_cache_chat(client, key, model, prefix)) + first = unwrap(send(prefix)) assert first.choices, f"{model}: first cache-priming call returned no choices: {first}" deadline = time.monotonic() + 30.0 while True: - second = unwrap(_cache_chat(client, key, model, prefix)) + second = unwrap(send(prefix)) read_tokens = _cached_read_tokens(second.usage) if read_tokens > 0 or time.monotonic() >= deadline: break @@ -125,7 +152,8 @@ class TestCacheControl: LiteLLMParamsBody(model=BEDROCK_MODEL, aws_region_name="us-east-1"), ) resources.defer(lambda: client.proxy.delete_model(model_id)) - _assert_cache_read_on_second_call(client, resources.key(), model) + key = resources.key() + _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix)) @pytest.mark.covers( "llm.chat_completions.vertex.prompt_cache_5m.nonstream.works", @@ -145,4 +173,40 @@ class TestCacheControl: ), ) resources.defer(lambda: client.proxy.delete_model(model_id)) - _assert_cache_read_on_second_call(client, resources.key(), model) + key = resources.key() + _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix)) + + @pytest.mark.covers( + "llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works", + exercised_on=[], + ) + def test_anthropic_prompt_caching_reads_cache( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-anthropic-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=ANTHROPIC_MODEL, api_key="os.environ/ANTHROPIC_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix)) + + @pytest.mark.covers( + "llm.chat_completions.openai.prompt_cache_5m.nonstream.works", + exercised_on=[], + ) + def test_openai_prompt_caching_reads_cache( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = f"e2e-openai-cache-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=OPENAI_MODEL, api_key="os.environ/OPENAI_API_KEY"), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + key = resources.key() + cache_key = f"e2e-openai-cache-{unique_marker()}" + _assert_cache_read_on_second_call( + model, lambda prefix: _plain_cache_chat(client, key, model, prefix, cache_key) + ) diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index f951eb328f5..5520ca0cee5 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -1,4 +1,4 @@ -"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex. +"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex, Cohere. Each test registers the deployment it needs at runtime (deleted on teardown) and asserts a non-empty, non-zero vector came back. The LIT-3167 guard in @@ -86,6 +86,26 @@ class TestEmbeddingsEndpoint: f"embedding vector is all zeros: {result.body[:300]}" ) + @pytest.mark.covers("llm.embeddings.cohere.basic.nonstream.works") + def test_cohere_embeddings_returns_vector( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = f"e2e-embeddings-cohere-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody(model="cohere/embed-v4.0", api_key="os.environ/COHERE_API_KEY"), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + result = endpoints_client.embeddings(key, model, "Say this is a test!") + require_successful_call(result) + parsed = EmbeddingsResult.model_validate_json(result.body) + assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}" + assert any(component != 0.0 for component in parsed.first_vector), ( + f"embedding vector is all zeros: {result.body[:300]}" + ) + @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works") def test_vertex_embeddings_returns_vector( self, endpoints_client: EndpointsClient, resources: ResourceManager diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index 50ea8f4b4df..e50e83eaf77 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -17,7 +17,7 @@ import pytest from e2e_config import CHEAP_OPENAI_MODEL, unique_marker from e2e_http import require_successful_call, unwrap from lifecycle import ResourceManager -from models import KeyGenerateBody, SpendLogRow +from models import ChatResponse, KeyGenerateBody, SpendLogRow from passthrough_client import ( AnthropicTool, GeminiFunctionDeclaration, @@ -344,6 +344,38 @@ class TestOpenAIPassthroughSpend: ) +class TestOpenAIProviderPrefixChat: + """OpenAI-format chat through the raw `/openai/{endpoint}` passthrough (LIT-4752). + + The body goes to OpenAI untranslated with the proxy's own OPENAI_API_KEY swapped + in, so the customer gets OpenAI's real completion back, and the gateway must + still write a costed pass_through_endpoint row for it. + """ + + @pytest.mark.covers("llm.chat_completions.openai.passthrough.nonstream.cost_logged") + def test_openai_prefix_chat_returns_completion_and_logs_its_cost( + self, client: PassthroughClient, scoped_key: str + ) -> None: + result = client.openai_chat(scoped_key, CHEAP_OPENAI_MODEL, f"Say hi in one word. {unique_marker()}") + require_successful_call(result) + + completion = ChatResponse.model_validate_json(result.body) + assert completion.id, f"/openai/v1/chat/completions relayed no completion id: {result.body[:300]}" + content = completion.choices[0].message.content if completion.choices and completion.choices[0].message else None + assert content and content.strip(), f"/openai/v1/chat/completions relayed an empty completion: {result.body[:300]}" + assert completion.usage is not None, f"the completion carried no usage to price from: {completion}" + + row = _fetch_cost_breakdown(client, completion.id) + assert row.prompt_tokens == completion.usage.prompt_tokens, ( + f"logged {row.prompt_tokens} prompt tokens, the completion the customer read " + f"reported {completion.usage.prompt_tokens}" + ) + assert row.completion_tokens == completion.usage.completion_tokens, ( + f"logged {row.completion_tokens} completion tokens, the completion the customer read " + f"reported {completion.usage.completion_tokens}" + ) + + class TestOpenAIPassthroughWebsocket: """The OpenAI passthrough prefixes must answer a websocket upgrade, not only a POST. From ff942c3a74e3d5a40e44f97111fffae56f29c06e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 11:05:50 -0700 Subject: [PATCH 39/61] fix(auto-router compression): surface a stored model-only policy in the edit form The backend treats either compression key on its own as an authoritative policy, but hydrate returned the untouched inherit state whenever the routing key was absent. A config carrying only auto_router_model_compression was therefore invisible in the form, and picking a routing value then overwrote the stored model hop. Only neither key set now reads as untouched, and an absent key on either hop hydrates as no compression for that hop rather than same-as-the-other. --- .../buildAutoRouterCompression.test.ts | 15 ++++++++++++++ .../add_model/buildAutoRouterCompression.ts | 20 ++++++++++++------- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts index ea6eaf99106..20d6af50d18 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.test.ts @@ -96,6 +96,21 @@ describe("hydrateAutoRouterCompression", () => { expect(rebuilt.auto_router_model_compression).not.toBe("headroom-a"); }); + it("surfaces a stored model-only policy instead of reading as untouched", () => { + // Regression: the backend treats either key alone as an authoritative policy, so a + // model-only config that hydrated to the inherit state was invisible in the form, + // and the next save overwrote the stored model hop with the routing value. + const state = hydrateAutoRouterCompression({ auto_router_model_compression: "headroom-b" }); + expect(state).toEqual({ routing: "none", sameAsRouting: false, model: "headroom-b" }); + }); + + it("round-trips a model-only policy without changing either hop", () => { + const stored = { auto_router_model_compression: "headroom-b" }; + const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored)); + expect(rebuilt.auto_router_model_compression).toBe("headroom-b"); + expect(rebuilt.auto_router_routing_compression).toBe("none"); + }); + it("round-trips through buildAutoRouterCompressionParams", () => { const original = { auto_router_routing_compression: "headroom-a", auto_router_model_compression: "none" }; const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(original)); diff --git a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts index 47d0e3db7e4..6f401a12865 100644 --- a/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts +++ b/ui/litellm-dashboard/src/components/add_model/buildAutoRouterCompression.ts @@ -50,14 +50,20 @@ export const hydrateAutoRouterCompression = (litellmParams: { auto_router_routing_compression?: string | null; auto_router_model_compression?: string | null; }): AutoRouterCompressionState => { - const routing = litellmParams.auto_router_routing_compression ?? undefined; - if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + const storedRouting = litellmParams.auto_router_routing_compression ?? undefined; + const storedModel = litellmParams.auto_router_model_compression ?? undefined; - // An absent model key is no model-hop compression, not same-as-routing: the backend - // reads it as None (policy_from_litellm_params). Hydrating it as same-as-routing - // would make re-saving an unrelated edit write the routing guardrail onto the model - // hop and silently start compressing the model call. - const model = litellmParams.auto_router_model_compression ?? NO_COMPRESSION; + // Only neither key set means the section was never touched. The backend treats + // either key on its own as an authoritative policy (policy_from_litellm_params), so + // reading a model-only config as untouched would hide it from the form and let the + // next save overwrite the stored model hop. + if (storedRouting === undefined && storedModel === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION; + + // An absent key on either hop is no compression for that hop, not same-as-the-other: + // the backend reads it as None. Hydrating it as same-as-routing would make re-saving + // an unrelated edit write one hop's guardrail onto the other. + const routing = storedRouting ?? NO_COMPRESSION; + const model = storedModel ?? NO_COMPRESSION; const sameAsRouting = model === routing; return { routing, sameAsRouting, model: sameAsRouting ? undefined : model }; }; From 723bc2140fb55dc7f7fdf56cac184763e092c8ae Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 11:19:39 -0700 Subject: [PATCH 40/61] refactor(auto-router compression): resolve the policy without a loop-local rebind The marker walk rebound a loop-local on each iteration, which is the mutation the repository's convention exists to discourage, but a `: Final` cannot express that inside a loop body: basedpyright rejects it outright with 'A Final variable cannot be assigned within a loop'. A lazy generator binds the name once per item and never rebinds it, so the first marker carrying a policy still wins and the rest are never read. --- litellm/proxy/guardrails/auto_router_compression.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index b6f32c1c46d..ab3ba4011c7 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -117,11 +117,11 @@ def policy_for_model( # request does not carry describes a different slice of traffic, so falling back # to it would apply, say, an "eu" policy to a "us" request purely on config order. untagged: Final = tuple(params for params in markers if not params.get("tags")) - for params in (*tag_matched, *untagged): - policy = policy_from_litellm_params(params) - if policy is not None: - return policy - return None + # Lazily, so the first marker carrying a policy still wins and the rest are never + # read. A generator rather than a loop-local: the name is bound once per item and + # never rebound, which `: Final` cannot express inside a loop body. + candidates: Final = (policy_from_litellm_params(params) for params in (*tag_matched, *untagged)) + return next((policy for policy in candidates if policy is not None), None) def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: From a0b07b47912caf10488f0e7926807cc83f5611d7 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 5 Sep 2026 11:33:21 -0700 Subject: [PATCH 41/61] docs(auto-router compression): cut the explanatory comments back The module, its routing hook and its tests carried long prose rationale where the repository allows only concise comments for genuinely complex logic. Trimmed to the non-obvious reasons and dropped the rest; no logic or test behaviour changes. --- .../guardrails/auto_router_compression.py | 89 ++++++------------- litellm/router.py | 27 ++---- .../test_auto_router_compression.py | 59 ++++-------- 3 files changed, 49 insertions(+), 126 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index ab3ba4011c7..98707e7ddca 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -1,15 +1,10 @@ """ -Decouples prompt compression between an auto router's routing decision and the -model it routes to. An auto router marker deployment may set -``auto_router_routing_compression`` and/or ``auto_router_model_compression`` in its -``litellm_params`` to name the compression guardrail that hop should use, or -``"none"`` to run no compression on that hop. Neither key set means the request's -own compression guardrails (key/team/model-level, or an "Always on" guardrail) -apply to both hops unchanged, exactly as before this feature existed. +Decouples prompt compression between an auto router's routing decision and the model +it routes to, via ``auto_router_routing_compression`` / ``auto_router_model_compression`` +on the marker deployment: a guardrail name, or ``"none"``. -Once either key is set, this auto router is authoritative: every other compression -guardrail is suppressed for that request, and only these two settings decide what -each hop sees. +Neither key set inherits today's behaviour. Either key set makes the auto router +authoritative and suppresses every other compression guardrail for that request. """ import contextvars @@ -29,12 +24,8 @@ if TYPE_CHECKING: COMPRESSION_GUARDRAIL_PROVIDERS: Final = frozenset({"headroom", "compresr"}) _NO_COMPRESSION: Final = "none" -# Compression guardrails this request's auto router has switched off. Deliberately a -# ContextVar rather than a metadata key: `refresh_proxy_server_request_body_snapshot` -# copies metadata into `proxy_server_request.body`, which deployments persist to spend -# logs. A suppression list that reaches a log the caller can read is a list the caller -# can replay, which would let any request switch off a PII or content-filter guardrail. -# Nothing here is caller-supplied, so there is no marker to forge in the first place. +# A ContextVar, not metadata: metadata reaches spend logs the caller can read, and a +# suppression list they can read is one they can replay to disable any guardrail. _suppressed_compression_guardrails: Final[contextvars.ContextVar[frozenset[str]]] = contextvars.ContextVar( "litellm_auto_router_suppressed_compression_guardrails", default=frozenset() ) @@ -45,9 +36,8 @@ def suppressed_compression_guardrails() -> frozenset[str]: return _suppressed_compression_guardrails.get() -# Whether `arm_pre_call` actually armed a model-side compression guardrail for this -# request. Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and -# nothing compresses; the router must not assume the model hop already ran. +# Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and nothing +# compresses; the router must not assume the model hop already ran. _model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( "litellm_auto_router_model_hop_armed", default=False ) @@ -95,10 +85,8 @@ def policy_for_model( ) -> AutoRouterCompressionPolicy | None: """The compression policy of the auto router marker `model_alias` resolves to. - Both the proxy's pre-call arming and the router's routing hook resolve the policy - through here, with the same tag rule, so an alias carrying several tag-scoped - markers can never suppress one marker's guardrail and then route under another - marker's policy. + Pre-call arming and the routing hook both resolve through here, so an alias with + several tag-scoped markers cannot suppress under one and then route under another. """ if llm_router is None: return None @@ -113,13 +101,9 @@ def policy_for_model( tag_matched: Final = tuple( params for params in markers if (tags := params.get("tags")) and requested.issuperset(frozenset(tags)) ) - # Only untagged markers may serve as the fallback. A marker scoped to tags this - # request does not carry describes a different slice of traffic, so falling back - # to it would apply, say, an "eu" policy to a "us" request purely on config order. + # Untagged only: a marker scoped to tags this request lacks describes other traffic. untagged: Final = tuple(params for params in markers if not params.get("tags")) - # Lazily, so the first marker carrying a policy still wins and the rest are never - # read. A generator rather than a loop-local: the name is bound once per item and - # never rebound, which `: Final` cannot express inside a loop body. + # Lazy, so the first marker carrying a policy wins and the rest are never read. candidates: Final = (policy_from_litellm_params(params) for params in (*tag_matched, *untagged)) return next((policy for policy in candidates if policy is not None), None) @@ -145,12 +129,8 @@ def _compression_guardrail_classes() -> tuple[type, ...]: def is_compression_guardrail(guardrail: object) -> bool: """Whether `guardrail` is an instance of a compression guardrail provider. - Both hops are validated through here. The two policy fields are operator-supplied - names and nothing else constrains them, so without this a name that resolves to an - ordinary guardrail would be handed the conversation and invoked: the routing hop - calls `apply_guardrail` directly, which POSTs the content wherever that guardrail - sends it, and the model hop is added to `metadata["guardrails"]`, which runs it even - when it is not `default_on`. + Both hops validate through here: the policy fields are operator-supplied names, and + an unvalidated one would get handed the conversation and invoked. """ classes: Final = _compression_guardrail_classes() return bool(classes) and isinstance(guardrail, classes) @@ -185,9 +165,6 @@ async def arm_pre_call( if not isinstance(model_alias, str) or not model_alias: return - # Read-only until a policy is confirmed: creating the metadata bucket for every - # request, including the vast majority with no auto-router compression policy, - # would be an unwanted side effect of merely checking for one. from litellm.router_strategy.tag_based_routing import ( _get_tags_from_request_kwargs, # pyright: ignore[reportPrivateUsage] # used in router.py and budget_limiter.py too ) @@ -209,8 +186,7 @@ async def arm_pre_call( ) ) - # Only a name that resolves to a real compression guardrail may be armed: this adds - # it to `metadata["guardrails"]`, which runs it even when it is not `default_on`. + # Arming adds the name to `metadata["guardrails"]`, which runs it even if not default_on. armed_model_hop: Final = policy.model is not None and any( guardrail.guardrail_name == policy.model for guardrail in _active_compression_guardrails() ) @@ -226,8 +202,7 @@ async def arm_pre_call( requested: Final = metadata.get("guardrails") existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else () if policy.model not in existing: - # A list, not a tuple: litellm_pre_call_utils tests this key with - # isinstance(..., list) and extends it, and would drop a tuple on the floor. + # A list: litellm_pre_call_utils isinstance-checks this key and drops a tuple. metadata["guardrails"] = [*existing, policy.model] # mutable-ok: this key's contract is a list @@ -240,25 +215,17 @@ def _as_routing_messages( async def messages_for_routing( policy: AutoRouterCompressionPolicy | None, - # list[dict], not Sequence[Mapping]: the async_pre_routing_hook protocol in - # litellm/types/router.py types `messages` as list[dict[str, Any]]. + # list[dict], not Sequence[Mapping]: fixed by the async_pre_routing_hook protocol. messages: list[dict[str, object]] | None, # mutable-ok: shape fixed by the pre-routing hook protocol request_kwargs: Mapping[str, object], ) -> list[dict[str, object]] | None: # mutable-ok: shape fixed by the pre-routing hook protocol - """Messages to use for a routing decision, per `policy.routing`. + """Messages to use for a routing decision, per `policy.routing`. None means the + caller should route on whatever it already has. - Returns None when the caller should route on whatever messages it already has. - - Always reads the live messages, never a pre-guardrail copy of them. The routing - hop compresses through a real guardrail, which POSTs the text to an external - compression service, so it must see what every other guardrail has already done - to the request. Routing on a snapshot taken before the pre-call hook would send - a masking guardrail's own input straight back out of the proxy. - - The consequence, when the model hop compressed and the two hops differ: the - messages in hand are that guardrail's output, and there is no un-compressed copy - left to route on. The routing decision reads the compressed text in that one - combination rather than leaking the original. + Reads the live messages, never a pre-guardrail copy: this compresses through a real + guardrail that POSTs the text out, so routing on a pre-masking snapshot would leak + what the masking guardrail stripped. When the model hop already compressed and the + hops differ, routing therefore reads the compressed text rather than the original. """ if policy is None or policy.routing is None: return None @@ -277,9 +244,6 @@ async def messages_for_routing( ) return _as_routing_messages(messages) - # apply_guardrail below hands this guardrail the conversation and it POSTs the - # content to whatever service backs it, so the name has to be a compression - # guardrail rather than any guardrail the operator happened to name. if not is_compression_guardrail(guardrail): verbose_proxy_logger.warning( "AutoRouter compression: guardrail '%s' is not a compression guardrail; routing on uncompressed messages", @@ -291,9 +255,8 @@ async def messages_for_routing( "structured_messages": _as_routing_messages(messages) # pyright: ignore[reportAssignmentType] # plain dicts, not AllMessageValues; see headroom.py's own use of this shape } model: Final = request_kwargs.get("model") - # A throwaway request_data: apply_guardrail writes its stats onto this dict, not the - # real request's metadata, so routing-side compression never double-counts against - # extract_compression_saved_tokens's model-savings accounting. + # Throwaway: apply_guardrail writes stats here, so routing never double-counts into + # extract_compression_saved_tokens. stats_sink: Final = {"messages": messages, "model": model} # mutable-ok: apply_guardrail writes its stats here result: Final = await guardrail.apply_guardrail( inputs=inputs, diff --git a/litellm/router.py b/litellm/router.py index 81de4572af8..51e2dbe38e4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13047,25 +13047,17 @@ class Router: team_id_from_request, ) - # Resolved through the same tag-aware lookup the proxy's pre-call arming used, - # so an alias carrying several tag-scoped markers cannot suppress one marker's - # guardrail and then route under a different marker's policy. + # Same tag-aware lookup the proxy's pre-call arming used, so an alias with + # several tag-scoped markers cannot suppress under one and route under another. compression_policy: Final = policy_for_model( llm_router=self, model_alias=registered_model_name, team_id=team_id_from_request(request_kwargs), request_tags=_get_tags_from_request_kwargs(request_kwargs), ) - # When both hops share the same compression, the model-side guardrail already - # ran in the proxy's ordinary pre-call hook and compressed `messages` in place - # (arm_pre_call armed it whether or not it is `default_on`); reuse that result - # for routing too instead of paying for a second compression call against the - # same content. - # - # Only the proxy calls arm_pre_call, so that reuse is conditional on it having - # actually run: on the SDK path nothing arms the model hop and nothing has - # compressed anything, and taking the shortcut there would skip both hops and - # silently serve the request with no compression at all. + # Shared compression already ran in the pre-call hook, so reuse it rather than + # compressing twice. Conditional on arming having actually happened: only the + # proxy arms, and on the SDK path the shortcut would skip both hops entirely. needs_independent_routing_compression: Final = compression_policy is not None and not ( compression_policy.is_same and compression_policy.model is not None and model_hop_compression_armed() ) @@ -13082,12 +13074,9 @@ class Router: input=input, specific_deployment=specific_deployment, ) - # The strategy only echoes back whatever `messages` it was handed, so a - # routing-only compression must not leak into the response: the model call - # and downstream deployment-context filtering both key off this field. - # Compared by value, not identity: PreRoutingHookResponse is a pydantic model, - # and pydantic reconstructs a validated list field rather than keeping the - # exact object passed in, even when nothing about it changed. + # Routing-only compression must not leak into the response: the model call and + # deployment-context filtering key off this field. Compared by value, since + # pydantic rebuilds the list rather than keeping the object passed in. pre_routing_hook_response: Final = ( routed.model_copy(update={"messages": messages}) # mutable-ok: pydantic's model_copy takes a dict if routed is not None and routing_messages is not None and routed.messages == routing_messages diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index 676b3ba2967..db2f94306fb 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -1,20 +1,4 @@ -""" -Unit tests for litellm.proxy.guardrails.auto_router_compression. - -Covers: -- policy_from_litellm_params: absent keys mean no policy; the "none" sentinel - normalizes to explicit no-compression within an active policy; is_same -- policy_for_model: finds the auto-router marker deployment for an alias, picks the - tag-scoped marker the request's tags actually match, and never falls back to a - marker scoped to tags the request does not carry -- arm_pre_call: no-op without a policy; suppresses active compression guardrails - through request-scoped state rather than metadata, which reaches spend logs a - caller can read; arms the model-side guardrail even when it isn't default_on -- messages_for_routing: no-op without a policy; compresses the live messages every - earlier guardrail has already rewritten, never a pre-guardrail copy of them; - never writes stats onto the caller's own request_kwargs (regression for - double-counted compression savings) -""" +"""Unit tests for litellm.proxy.guardrails.auto_router_compression.""" import json from typing import Any @@ -137,8 +121,7 @@ class TestPolicyForModel: assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) def test_a_marker_scoped_to_other_tags_is_never_the_fallback(self): - """Regression: an "eu" marker describes a different slice of traffic, so a "us" - request must not fall back to its policy just because it is configured first.""" + """Regression: a "us" request must not fall back to an "eu" marker's policy.""" router = _FakeRouter( [ _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), @@ -149,8 +132,7 @@ class TestPolicyForModel: assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) def test_no_untagged_fallback_means_no_policy(self): - """With only tag-scoped markers and none matching, there is no policy to apply: - inheriting an unrelated slice's compression is worse than inheriting nothing.""" + """No matching marker means no policy, not an unrelated slice's compression.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) assert ( policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None @@ -268,9 +250,8 @@ class TestArmPreCall: @pytest.mark.asyncio async def test_suppression_state_never_enters_request_metadata(self): - """Regression (security): a suppression list written to metadata is copied into - proxy_server_request.body and persisted to spend logs, so a caller could read it - back and replay it to switch off a PII or content-filter guardrail.""" + """Regression (security): metadata reaches spend logs, so a suppression list + there is one a caller could read back and replay to disable a guardrail.""" guardrail = _RecordingCompressionGuardrail(guardrail_name="always-on-compression") import litellm @@ -312,9 +293,8 @@ class TestArmPreCall: @pytest.mark.asyncio async def test_arm_pre_call_keeps_no_copy_of_the_prompt(self): - """Regression (security): arm_pre_call runs before the pre-call guardrails, so - any copy of the messages it retained would be the pre-masking text. Routing-side - compression POSTs its input to an external service, so that copy must not exist.""" + """Regression (security): arm_pre_call runs before the guardrails, so any copy it + kept would be pre-masking text that routing then POSTs to an external service.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) data = {"model": "smart-router", "messages": [{"role": "user", "content": "my ssn is 123-45-6789"}]} @@ -337,10 +317,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_routing_none_never_reaches_for_a_pre_guardrail_copy(self): - """Routing asked for no compression while the model hop compressed, so the - messages in hand are that guardrail's output and no uncompressed copy survives. - Routing reads them as-is: the alternative is keeping a pre-guardrail copy, which - is the text a masking guardrail exists to remove.""" + """No uncompressed copy survives the model hop, and keeping one would mean + retaining the pre-masking text. Routing reads what it has.""" policy = AutoRouterCompressionPolicy(routing=None, model="headroom-a") model_compressed = [{"role": "user", "content": "[COMPRESSED] the full original conversation"}] @@ -362,10 +340,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_routing_compresses_what_the_other_guardrails_left_behind(self, registered_guardrail): - """Regression (security): routing-side compression POSTs its input to an external - service, so it must read the live messages every earlier guardrail has already - rewritten. Routing on a pre-guardrail copy would send a masking guardrail's own - input straight back out of the proxy.""" + """Regression (security): routing POSTs its input out, so it must read what the + earlier guardrails left behind, not a pre-masking copy.""" policy = AutoRouterCompressionPolicy(routing="fake-compress", model="headroom-b") masked = [{"role": "user", "content": "my ssn is [REDACTED]"}] @@ -376,10 +352,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_a_non_compression_guardrail_is_never_invoked_for_routing(self, monkeypatch): - """Regression (security): the policy fields are operator-supplied names that - nothing else constrains. apply_guardrail hands the guardrail the conversation - and it POSTs that content to whatever service backs it, so naming an ordinary - guardrail must not turn the routing hop into a way to ship prompts there.""" + """Regression (security): naming an ordinary guardrail must not turn the routing + hop into a way to ship prompts to whatever service backs it.""" import litellm other = _NonCompressionGuardrail(guardrail_name="pii-filter") @@ -397,11 +371,8 @@ class TestMessagesForRouting: @pytest.mark.asyncio async def test_guardrail_receives_a_throwaway_request_data_not_the_real_request_kwargs(self, registered_guardrail): - """Regression: a real compression guardrail writes its stats onto whatever - `request_data` dict it's given (`add_standard_logging_guardrail_information_to_ - request_data`). If that were the caller's own `request_kwargs`, routing-side - compression would double-count into extract_compression_saved_tokens, which - sums every guardrail_information entry on the real request's metadata.""" + """Regression: a guardrail writes stats onto the request_data it is given, so + passing the caller's own would double-count into extract_compression_saved_tokens.""" policy = AutoRouterCompressionPolicy(routing="fake-compress", model=None) messages = [{"role": "user", "content": "hi"}] request_kwargs = {"metadata": {}} From 4df284e16dcc02ceabad4576df6f2c976f20d839 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:39:24 -0700 Subject: [PATCH 42/61] fix(guardrails): record guardrail information for undecorated custom apply_guardrail overrides (#39727) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 3 + litellm/integrations/custom_guardrail.py | 9 + .../integrations/test_custom_guardrail.py | 168 ++++++++++++++++++ .../test_openai_guardrail_handler.py | 31 ++++ .../test_bedrock_guardrails.py | 6 +- 5 files changed, 216 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 7d6de612349..e4fb2a00297 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -215,6 +215,9 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # so the deployment-level hook does not re-run them for the same request PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" +# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again +LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information" + # Generic fallback for unknown models DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9bb613654e4..c462b1edb98 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -46,6 +46,7 @@ dc: Final = DualCache() from litellm.constants import ( GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS, + LOGS_GUARDRAIL_INFORMATION_MARKER, PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) from litellm.exceptions import ( @@ -151,6 +152,13 @@ class CustomGuardrail(CustomLogger): records_own_guardrail_information: ClassVar[bool] = False + def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks + super().__init_subclass__(**kwargs) + own_apply_guardrail: Final = cls.__dict__.get("apply_guardrail") + if own_apply_guardrail is None or LOGS_GUARDRAIL_INFORMATION_MARKER in vars(own_apply_guardrail): + return + cls.apply_guardrail = log_guardrail_information(own_apply_guardrail) + def __init__( self, guardrail_name: str | None = None, @@ -1559,4 +1567,5 @@ def log_guardrail_information(func): return async_wrapper(*args, **kwargs) return sync_wrapper(*args, **kwargs) + vars(wrapper)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the wrapper this call just built return wrapper diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 6edc2c9bf77..49a52157c8a 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,5 @@ import asyncio +from typing import TYPE_CHECKING, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -11,6 +12,9 @@ from litellm.integrations.custom_guardrail import ( from litellm.proxy._types import CallTypes, UserAPIKeyAuth from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class TestCustomGuardrailDeploymentHook: @@ -2239,6 +2243,170 @@ class TestRecordsOwnGuardrailInformation: assert _guardrail_entries(request_data) == [] +class _UndecoratedGuardrail(CustomGuardrail): + """apply_guardrail written like the docs example: no @log_guardrail_information.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + from litellm.exceptions import GuardrailRaisedException + + if any("forbidden" in text for text in inputs.get("texts") or []): + raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message="Content blocked") + return inputs + + +class _UndecoratedSelfRecordingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"custom": True}, + request_data=request_data, + guardrail_status="success", + start_time=0.0, + end_time=0.0, + duration=0.0, + ) + return inputs + + +class _InheritedApplyGuardrail(_UndecoratedGuardrail): + pass + + +class TestUndecoratedApplyGuardrailIsLogged: + """LIT-5983 regression: a custom guardrail that overrides apply_guardrail without the + @log_guardrail_information decorator must still record guardrail information, and the + auto-wrap must not double-record decorated or self-recording implementations.""" + + @pytest.mark.asyncio + async def test_undecorated_success_is_recorded(self): + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = _UndecoratedGuardrail(guardrail_name="docs-style", event_hook=GuardrailEventHooks.pre_call) + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "docs-style" + assert entries[0]["guardrail_mode"] == "pre_call" + assert entries[0]["guardrail_status"] == "success" + + @pytest.mark.asyncio + async def test_undecorated_block_is_recorded_and_reraised(self): + from litellm.exceptions import GuardrailRaisedException + + guardrail = _UndecoratedGuardrail(guardrail_name="docs-style") + request_data: dict = {"model": "gpt-4o"} + + with pytest.raises(GuardrailRaisedException): + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["forbidden"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "docs-style" + assert entries[0]["guardrail_status"] == "guardrail_intervened" + + @pytest.mark.asyncio + async def test_undecorated_bare_exception_is_recorded_as_failed_to_respond(self): + class _BareExceptionGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + raise Exception("Content blocked: policy violation") + + guardrail = _BareExceptionGuardrail(guardrail_name="docs-style") + request_data: dict = {"model": "gpt-4o"} + + with pytest.raises(Exception, match="Content blocked"): + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["x"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" + + @pytest.mark.asyncio + async def test_inherited_apply_guardrail_is_recorded_once(self): + guardrail = _InheritedApplyGuardrail(guardrail_name="child") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + assert len(_guardrail_entries(request_data)) == 1 + + @pytest.mark.asyncio + async def test_undecorated_self_recording_apply_guardrail_is_recorded_once(self): + guardrail = _UndecoratedSelfRecordingGuardrail(guardrail_name="self-recording") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + entries = _guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_response"] == {"custom": True} + + @pytest.mark.asyncio + async def test_base_apply_guardrail_is_not_recorded(self): + guardrail = CustomGuardrail(guardrail_name="base") + request_data: dict = {"model": "gpt-4o"} + + await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["hello"]), + request_data=request_data, + input_type="request", + ) + + assert _guardrail_entries(request_data) == [] + + def test_subclass_keywords_reach_cooperative_init_subclass(self): + class _LabelMixin: + seen_label: str = "" + + def __init_subclass__(cls, label: str = "", **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + cls.seen_label = label + + class _Labelled(CustomGuardrail, _LabelMixin, label="docs-style"): + pass + + assert _Labelled.seen_label == "docs-style" + + class _ApplyOnlyObserver(CustomGuardrail): """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index cebab2512d0..36e715d5804 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1075,6 +1075,37 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert result == responses_so_far +class TestUndecoratedGuardrailIsRecorded: + """LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail + without @log_guardrail_information must still end up in the request's guardrail + information on both the request and response paths.""" + + @pytest.mark.asyncio + async def test_request_path_records_undecorated_guardrail(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="docs-style") + data = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}} + + await handler.process_input_messages(data, guardrail) + + entries = data["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")] + + @pytest.mark.asyncio + async def test_response_path_records_undecorated_guardrail(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="docs-style") + response = ModelResponse( + choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))] + ) + request_data: dict = {"metadata": {}} + + await handler.process_output_response(response, guardrail, request_data=request_data) + + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")] + + class TestGetStructuredMessages: """Test the get_structured_messages method.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 9842d88e8d1..7da55f22bda 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -1137,7 +1137,11 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): mock_api.assert_called_once() kwargs = mock_api.call_args.kwargs assert kwargs["source"] == "OUTPUT" - assert kwargs["request_data"] == {"model": "gpt-4o"} + assert kwargs["request_data"]["model"] == "gpt-4o" + recorded = kwargs["request_data"]["metadata"]["standard_logging_guardrail_information"] + assert [(e["guardrail_name"], e["guardrail_status"]) for e in recorded] == [ + (guardrail.guardrail_name, "success") + ] synthetic = kwargs["response"] assert isinstance(synthetic, ModelResponse) assert len(synthetic.choices) == 2 From f66b3ebe0dda8e25c47739c1eb15637dce2d11ad Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:40:00 -0700 Subject: [PATCH 43/61] feat(responses): honor supported_endpoints /v1/responses opt-in for OpenAI-compatible deployments (#39725) * feat(responses): honor supported_endpoints /v1/responses opt-in for OpenAI-compatible deployments custom_openai and other generic OpenAI-compatible deployments have no native Responses API config, so every /v1/responses call is bridged through /v1/chat/completions. When model_info.supported_endpoints lists /v1/responses, resolve OpenAILikeResponsesConfig instead so the request is forwarded to {api_base}/responses, for streaming, non-streaming and mode: responses deployments alike. Providers with their own Responses config are unchanged. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): drop deployment supported_endpoints opt-in after cross-provider prompt swap A prompt manager that moves the request to another provider leaves kwargs['model_info'] describing the original deployment; without this the swapped provider was sent an OpenAI-like /responses request it does not serve. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(responses): carry prompt-swap deployment metadata as a return value instead of a kwargs marker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/main.py | 93 +++++-- ...sponses_supported_endpoints_passthrough.py | 254 ++++++++++++++++++ 2 files changed, 327 insertions(+), 20 deletions(-) create mode 100644 tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py diff --git a/litellm/responses/main.py b/litellm/responses/main.py index ed2d6a216fd..5e74b7324b4 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2,6 +2,7 @@ import asyncio import contextvars from collections.abc import Coroutine, Generator, Iterable, Mapping from contextlib import contextmanager +from dataclasses import dataclass from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast @@ -23,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) @@ -403,8 +405,40 @@ def _bridges_to_chat_completions( return responses_api_provider_config is None or use_chat_completions_api is True +def _deployment_passes_through_responses(model_info: object) -> bool: + """Whether ``model_info.supported_endpoints`` opts the deployment into native ``{api_base}/responses``.""" + if not isinstance(model_info, dict): + return False + supported_endpoints: Final = model_info.get("supported_endpoints") + return isinstance(supported_endpoints, (list, tuple)) and "/v1/responses" in supported_endpoints + + +def _deployment_model_info_after_prompt_swap( + requested_provider: str | None, resolved_provider: str | None, model_info: object +) -> object: + """Deployment metadata only describes the upstream while the prompt manager keeps its provider.""" + return model_info if resolved_provider == requested_provider else None + + +@dataclass(frozen=True, slots=True) +class _AsyncPromptManagementOutcome: + merged_optional_params: Mapping[str, object] + deployment_model_info: object + + +def _resolve_responses_api_provider_config( + model: str, custom_llm_provider: str, model_info: object +) -> BaseResponsesAPIConfig | None: + provider_config: Final = ProviderConfigManager.get_provider_responses_api_config( + model=model, provider=custom_llm_provider + ) + if provider_config is not None or not _deployment_passes_through_responses(model_info): + return provider_config + return OpenAILikeResponsesConfig() + + def _will_bridge_to_chat_completions( - model: str, custom_llm_provider: str | None, use_chat_completions_api: bool + model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object ) -> bool: """``_bridges_to_chat_completions`` for callers running before the provider config is resolved. @@ -418,9 +452,7 @@ def _will_bridge_to_chat_completions( if custom_llm_provider is None: return True return _bridges_to_chat_completions( - ProviderConfigManager.get_provider_responses_api_config( - model=normalized_model[0], provider=custom_llm_provider - ), + _resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info), use_chat_completions_api or normalized_model[1], ) @@ -527,7 +559,10 @@ async def aresponses( with _prompt_management_sees_a_provisional_message_list( kwargs, bridged=_will_bridge_to_chat_completions( - model, custom_llm_provider, bool(kwargs.get("use_chat_completions_api")) + model, + custom_llm_provider, + bool(kwargs.get("use_chat_completions_api")), + kwargs.get("model_info"), ), ): ( @@ -552,6 +587,7 @@ async def aresponses( merged_input=merged_input, ), ) + requested_provider: Final = custom_llm_provider if model != original_model: custom_llm_provider = _resolve_prompt_swapped_provider( original_model=original_model, @@ -561,7 +597,12 @@ async def aresponses( prompt_id=prompt_id, ) kwargs.pop("prompt_id", None) - kwargs["_async_prompt_merged_params"] = merged_optional_params + kwargs["_async_prompt_merged_params"] = _AsyncPromptManagementOutcome( + merged_optional_params=merged_optional_params, + deployment_model_info=_deployment_model_info_after_prompt_swap( + requested_provider, custom_llm_provider, kwargs.get("model_info") + ), + ) func: Final = partial( responses, @@ -666,12 +707,14 @@ def _apply_prompt_management_to_responses_call( kwargs: dict[str, Any], local_vars: dict[str, object], use_chat_completions_api: bool, -) -> tuple[str | ResponseInputParam, str, str | None]: - async_merged: Final[Mapping[str, object] | None] = kwargs.pop("_async_prompt_merged_params", None) - if async_merged is not None: - for key, value in async_merged.items(): +) -> tuple[str | ResponseInputParam, str, str | None, object]: + """Returns the prompt-managed input, model and provider, plus the deployment metadata that still + describes the upstream (``None`` once the prompt manager moved the request to another provider).""" + async_outcome: Final[_AsyncPromptManagementOutcome | None] = kwargs.pop("_async_prompt_merged_params", None) + if async_outcome is not None: + for key, value in async_outcome.merged_optional_params.items(): local_vars[key] = value - return input, model, custom_llm_provider + return input, model, custom_llm_provider, async_outcome.deployment_model_info prompt_id: Final = cast(str | None, kwargs.get("prompt_id", None)) prompt_variables: Final = cast(dict | None, kwargs.get("prompt_variables", None)) @@ -684,7 +727,9 @@ def _apply_prompt_management_to_responses_call( ): with _prompt_management_sees_a_provisional_message_list( kwargs, - bridged=_will_bridge_to_chat_completions(model, custom_llm_provider, use_chat_completions_api), + bridged=_will_bridge_to_chat_completions( + model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info") + ), ): ( model, @@ -710,19 +755,28 @@ def _apply_prompt_management_to_responses_call( ) local_vars["input"] = input local_vars["model"] = model - if model != original_model: - custom_llm_provider = _resolve_prompt_swapped_provider( + resolved_provider: Final = ( + custom_llm_provider + if model == original_model + else _resolve_prompt_swapped_provider( original_model=original_model, swapped_model=model, custom_llm_provider=custom_llm_provider, kwargs=kwargs, prompt_id=prompt_id, ) - local_vars["custom_llm_provider"] = custom_llm_provider + ) + local_vars["custom_llm_provider"] = resolved_provider for key, value in merged_optional_params.items(): local_vars[key] = value + return ( + input, + model, + resolved_provider, + _deployment_model_info_after_prompt_swap(custom_llm_provider, resolved_provider, kwargs.get("model_info")), + ) - return input, model, custom_llm_provider + return input, model, custom_llm_provider, kwargs.get("model_info") # Opt-in via model id (mirrors the `responses/` prefix pattern on chat completions). @@ -1052,7 +1106,7 @@ def responses( ) local_vars["custom_llm_provider"] = custom_llm_provider - input, model, custom_llm_provider = _apply_prompt_management_to_responses_call( + input, model, custom_llm_provider, deployment_model_info = _apply_prompt_management_to_responses_call( input=input, model=model, custom_llm_provider=custom_llm_provider, @@ -1123,9 +1177,8 @@ def responses( if custom_llm_provider is None: responses_api_provider_config = None else: - responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config( - model=model, - provider=custom_llm_provider, + responses_api_provider_config = _resolve_responses_api_provider_config( + model, custom_llm_provider, deployment_model_info ) local_vars.update(kwargs) diff --git a/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py b/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py new file mode 100644 index 00000000000..7cd04b015f9 --- /dev/null +++ b/tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py @@ -0,0 +1,254 @@ +""" +A deployment with `model_info.supported_endpoints` containing `/v1/responses` forwards +`/v1/responses` natively to `{api_base}/responses`. Without it, generic OpenAI-compatible +providers such as `custom_openai` keep bridging through `/v1/chat/completions`. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +import respx + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig +from litellm.responses.main import _resolve_responses_api_provider_config +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ModelResponse + +API_BASE = "https://backend.example/v1" +RESPONSES_URL = f"{API_BASE}/responses" +CHAT_URL = f"{API_BASE}/chat/completions" +OPT_IN = {"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]} + +RESPONSES_BODY = { + "id": "resp_native", + "object": "response", + "created_at": 1741476542, + "status": "completed", + "model": "my-model", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "native", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} + +CHAT_BODY = { + "id": "chatcmpl_bridged", + "object": "chat.completion", + "created": 1741476542, + "model": "my-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "bridged"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + +SSE_BODY = ( + "event: response.created\n" + f"data: {json.dumps({'type': 'response.created', 'response': RESPONSES_BODY})}\n\n" + "event: response.completed\n" + f"data: {json.dumps({'type': 'response.completed', 'response': RESPONSES_BODY})}\n\n" +) + + +def _mock_backend(router: respx.MockRouter) -> tuple[respx.Route, respx.Route]: + responses_route = router.post(RESPONSES_URL).mock(return_value=httpx.Response(200, json=RESPONSES_BODY)) + chat_route = router.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY)) + return responses_route, chat_route + + +SWAPPED_MODEL = "deepseek/deepseek-chat" +SWAPPED_API_BASE = "https://api.deepseek.com/beta" + + +def _prompt_manager_swapping_to(model: str) -> MagicMock: + """A logging object whose prompt hook rewrites the request's model, as a prompt manager does.""" + prompt_return = (model, [{"role": "user", "content": "hi"}], {}) + logging_obj = MagicMock() + logging_obj.__class__ = LiteLLMLoggingObj + logging_obj.should_run_prompt_management_hooks.return_value = True + logging_obj.get_chat_completion_prompt.return_value = prompt_return + logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return) + logging_obj.model_call_details = {} + return logging_obj + + +def _mock_swap_targets(router: respx.MockRouter, monkeypatch) -> tuple[respx.Route, respx.Route]: + """The swapped provider's chat endpoint, plus the `/responses` it does not serve but a stale + opt-in would send to.""" + monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek") + swapped_chat_route = router.post(f"{SWAPPED_API_BASE}/chat/completions").mock( + return_value=httpx.Response(200, json=CHAT_BODY) + ) + stale_responses_route = router.post(f"{SWAPPED_API_BASE}/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_BODY) + ) + return swapped_chat_route, stale_responses_route + + +@pytest.fixture(autouse=True) +def _respx_interceptable_httpx_client(monkeypatch): + monkeypatch.setattr(litellm, "num_retries", 0) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.parametrize( + "model_info, expected_type", + [ + (OPT_IN, OpenAILikeResponsesConfig), + ({"supported_endpoints": ["/v1/chat/completions"]}, type(None)), + ({}, type(None)), + (None, type(None)), + ("/v1/responses", type(None)), + ], +) +def test_resolver_opt_in_gates_openai_like_config(model_info, expected_type): + config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info) + assert type(config) is expected_type + + +def test_resolver_keeps_native_provider_config(): + """`openai/` already routes /v1/responses natively; the opt-in must not swap its config.""" + config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN) + assert type(config) is OpenAIResponsesAPIConfig + + +@respx.mock +async def test_opt_in_forwards_responses_natively(): + responses_route, chat_route = _mock_backend(respx.mock) + + result = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + api_base=API_BASE, + api_key="sk-backend", + model_info=OPT_IN, + ) + + assert responses_route.call_count == 1 + assert chat_route.call_count == 0 + request = responses_route.calls.last.request + assert request.headers["authorization"] == "Bearer sk-backend" + assert json.loads(request.content)["input"] == "hi" + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "native" + + +@respx.mock +async def test_opt_in_forwards_streaming_responses_natively(monkeypatch): + """The router registers each deployment in `litellm.model_cost`; an unregistered model is + treated as non-streaming and would be faked, so mirror that registration here.""" + monkeypatch.setitem(litellm.model_cost, "custom_openai/my-model", {"litellm_provider": "custom_openai"}) + responses_route = respx.post(RESPONSES_URL).mock( + return_value=httpx.Response(200, text=SSE_BODY, headers={"content-type": "text/event-stream"}) + ) + chat_route = respx.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY)) + + stream = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + stream=True, + api_base=API_BASE, + api_key="sk-backend", + model_info=OPT_IN, + ) + events = [event async for event in stream] + + assert responses_route.call_count == 1 + assert chat_route.call_count == 0 + assert json.loads(responses_route.calls.last.request.content)["stream"] is True + assert [event.type for event in events] == ["response.created", "response.completed"] + + +@respx.mock +async def test_without_opt_in_still_bridges_through_chat_completions(): + responses_route, chat_route = _mock_backend(respx.mock) + + result = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + api_base=API_BASE, + api_key="sk-backend", + model_info={"supported_endpoints": ["/v1/chat/completions"]}, + ) + + assert chat_route.call_count == 1 + assert responses_route.call_count == 0 + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "bridged" + + +@respx.mock +async def test_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch): + """When a prompt manager moves the request to another provider, the original deployment's + `supported_endpoints` no longer describes the upstream, so the swapped provider bridges.""" + swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch) + + result = await litellm.aresponses( + model="custom_openai/my-model", + input="hi", + prompt_id="p1", + litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL), + model_info=OPT_IN, + ) + + assert swapped_chat_route.call_count == 1 + assert stale_responses_route.call_count == 0 + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "bridged" + + +@respx.mock +def test_sync_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch): + swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch) + + result = litellm.responses( + model="custom_openai/my-model", + input="hi", + prompt_id="p1", + litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL), + model_info=OPT_IN, + ) + + assert swapped_chat_route.call_count == 1 + assert stale_responses_route.call_count == 0 + assert isinstance(result, ResponsesAPIResponse) + assert result.output[0].content[0].text == "bridged" + + +@respx.mock +async def test_mode_responses_chat_completion_reaches_native_responses(monkeypatch): + """A `mode: responses` deployment bridges chat completions into the Responses API; with + the opt-in that inner call must reach `{api_base}/responses` instead of bouncing back + to `/chat/completions`.""" + responses_route, chat_route = _mock_backend(respx.mock) + monkeypatch.setitem( + litellm.model_cost, + "custom_openai/my-model", + {"mode": "responses", "litellm_provider": "custom_openai"}, + ) + + result = await litellm.acompletion( + model="custom_openai/my-model", + messages=[{"role": "user", "content": "hi"}], + api_base=API_BASE, + api_key="sk-backend", + model_info={"mode": "responses", **OPT_IN}, + ) + + assert responses_route.call_count == 1 + assert chat_route.call_count == 0 + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "native" From e6705510f82c0c70b274c922210bbdd8edab5379 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:41:35 -0700 Subject: [PATCH 44/61] fix(router): hold max_parallel_requests slot until streaming response is exhausted or closed (#39859) * fix(router): hold max_parallel_requests slot until streaming response is exhausted or closed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(router): normalize deployment_slot once to keep stream_with_fallbacks under the C901 ceiling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): close upstream stream before releasing max_parallel_requests slot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 108 ++++++++++++----------- tests/test_litellm/test_router.py | 141 ++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 51 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 6943eece90f..490836f5f0a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2597,14 +2597,20 @@ class Router: model_response: CustomStreamWrapper, messages: list[dict[str, str]], initial_kwargs: dict, + deployment_slot: contextlib.AsyncExitStack | None = None, ) -> CustomStreamWrapper: """ Helper to iterate over a streaming response. Catches errors for fallbacks using the router's fallback system + + `deployment_slot` holds the deployment's max_parallel_requests semaphore; it is + released when the stream is exhausted, closed, or falls back to another deployment """ from litellm.exceptions import MidStreamFallbackError + held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() + class FallbackStreamWrapper(CustomStreamWrapper): def __init__(self, async_generator: AsyncGenerator): # Copy attributes from the original model_response @@ -2628,12 +2634,26 @@ class Router: async def __anext__(self): return await self._async_generator.__anext__() + async def close_model_response() -> None: + if not hasattr(model_response, "aclose"): + return + try: + await model_response.aclose() + except BaseException as e: + verbose_router_logger.debug( + "stream_with_fallbacks: error closing model_response: %s", + e, + ) + async def stream_with_fallbacks(): fallback_response = None # Track for cleanup in finally try: async for item in model_response: yield item except MidStreamFallbackError as e: + with anyio.CancelScope(shield=True): + await close_model_response() + await held_slot.aclose() if not e.is_pre_first_chunk and ( e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) ): @@ -2707,14 +2727,8 @@ class Router: # (e.g. on client disconnect). # Shield from anyio cancellation so the awaits can complete. with anyio.CancelScope(shield=True): - if hasattr(model_response, "aclose"): - try: - await model_response.aclose() - except BaseException as e: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - e, - ) + await close_model_response() + await held_slot.aclose() if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() @@ -3379,61 +3393,53 @@ class Router: kwargs=kwargs, client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore): - async with rpm_semaphore: - """ - - Check rpm limits before making the call - - If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe) - """ - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, - logging_obj=logging_obj, - parent_otel_span=parent_otel_span, - ) - response = await _response - else: + async with contextlib.AsyncExitStack() as deployment_slot: + if isinstance(rpm_semaphore, asyncio.Semaphore): + await deployment_slot.enter_async_context(rpm_semaphore) await self.async_routing_strategy_pre_call_checks( deployment=deployment, logging_obj=logging_obj, parent_otel_span=parent_otel_span, ) - response = await _response - ## CHECK CONTENT FILTER ERROR ## - if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs) - if _should_raise: - raise litellm.ContentPolicyViolationError( - message="Response output was blocked.", - model=model, - llm_provider="", + ## CHECK CONTENT FILTER ERROR ## + if isinstance(response, ModelResponse): + _should_raise = self._should_raise_content_policy_error( + model=model, response=response, kwargs=kwargs ) + if _should_raise: + raise litellm.ContentPolicyViolationError( + message="Response output was blocked.", + model=model, + llm_provider="", + ) - if ( - isinstance(response, CustomStreamWrapper) - and response.completion_stream is None - and response.make_call is not None - ): - await response.fetch_stream() + if ( + isinstance(response, CustomStreamWrapper) + and response.completion_stream is None + and response.make_call is not None + ): + await response.fetch_stream() - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - # debug how often this deployment picked - self._track_deployment_metrics( - deployment=deployment, - response=response, - parent_otel_span=parent_otel_span, - ) - - if isinstance(response, CustomStreamWrapper): - return await self._acompletion_streaming_iterator( - model_response=response, - messages=messages, - initial_kwargs=input_kwargs_for_streaming_fallback, + self.success_calls[model_name] += 1 + verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + # debug how often this deployment picked + self._track_deployment_metrics( + deployment=deployment, + response=response, + parent_otel_span=parent_otel_span, ) - return response + if isinstance(response, CustomStreamWrapper): + return await self._acompletion_streaming_iterator( + model_response=response, + messages=messages, + initial_kwargs=input_kwargs_for_streaming_fallback, + deployment_slot=deployment_slot.pop_all(), + ) + + return response except litellm.Timeout as e: deployment_request_timeout_param: Final = _timeout_debug_deployment_dict.get("litellm_params", {}).get( "request_timeout", None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 31eb46f1458..ffd6c5f97ce 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12938,3 +12938,144 @@ async def test_router_retry_policy_controls_upstream_attempt_count( await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) assert upstream.call_count == expected_upstream_calls + + +class _InFlightTracker: + def __init__(self) -> None: + self.current = 0 + self.peak = 0 + + def enter(self) -> None: + self.current += 1 + self.peak = max(self.peak, self.current) + + def exit(self) -> None: + self.current -= 1 + + +_SSE_CHUNKS: Final[tuple[bytes, ...]] = tuple( + b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",' + b'"choices":[{"index":0,"delta":{"content":"x"},"finish_reason":null}]}\n\n' + for _ in range(5) +) + + +class _CountingSSEStream(httpx.AsyncByteStream): + def __init__(self, tracker: _InFlightTracker) -> None: + self._tracker = tracker + self._in_flight = False + + def _finish(self) -> None: + if self._in_flight: + self._in_flight = False + self._tracker.exit() + + async def __aiter__(self): + self._in_flight = True + self._tracker.enter() + try: + for chunk in _SSE_CHUNKS: + await asyncio.sleep(0.02) + yield chunk + finally: + await self.aclose() + yield b"data: [DONE]\n\n" + + async def aclose(self) -> None: + await asyncio.sleep(0.02) + self._finish() + + +def _max_parallel_router(max_parallel_requests: int) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-5.6", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "sk-fake", + "api_base": "https://max-parallel.local/v1", + "max_parallel_requests": max_parallel_requests, + }, + } + ], + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_router_max_parallel_requests_bounds_in_flight_upstream_calls( + monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + tracker: Final = _InFlightTracker() + router: Final = _max_parallel_router(max_parallel_requests=2) + + async def upstream(request: httpx.Request) -> httpx.Response: + if stream: + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker) + ) + tracker.enter() + await asyncio.sleep(0.05) + tracker.exit() + return httpx.Response( + 200, + json={ + "id": "c", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}], + }, + ) + + async def one_call() -> None: + response = await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=stream + ) + if stream: + async for _ in response: + pass + + with respx.mock(assert_all_called=True) as respx_mock: + respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream) + await asyncio.wait_for(asyncio.gather(*(one_call() for _ in range(10))), timeout=10) + + assert tracker.peak <= 2 + assert tracker.current == 0 + + +@pytest.mark.asyncio +async def test_router_max_parallel_requests_slot_released_when_stream_closed_early(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + tracker: Final = _InFlightTracker() + router: Final = _max_parallel_router(max_parallel_requests=1) + + with respx.mock() as respx_mock: + respx_mock.post("https://max-parallel.local/v1/chat/completions").mock( + side_effect=lambda request: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker) + ) + ) + first: Final = await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True + ) + await first.__anext__() + + async def second_call() -> None: + second = await router.acompletion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True + ) + async for _ in second: + pass + + second_task: Final = asyncio.create_task(second_call()) + await asyncio.sleep(0.05) + assert tracker.current == 1 + await first.aclose() + await asyncio.wait_for(second_task, timeout=2) + + assert tracker.peak == 1 + assert tracker.current == 0 From cba3dd58287114588dad0624cd2c8d0d040b902d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:41:56 -0700 Subject: [PATCH 45/61] fix(proxy): retry deadlocks and requeue spend logs on any DB write error (#39883) * fix(proxy): retry deadlocks and requeue spend logs on any DB write error update_spend_logs dequeued the batch and only retried/requeued on transport errors. A 40P01 deadlock surfaced as a plain prisma DataError and went through poison-row isolation, which dropped every row it hit; every other DB error was re-raised with the batch already gone from the queue. Treat deadlocks as transient (retry, then requeue), keep them out of poison-row isolation, and requeue the batch at the head of the queue on any other prisma error so it lands once the DB is healthy. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): drop redundant docstrings and tighten test typing for spend-log requeue Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert deadlock retries from mock call history instead of mutable lists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/exception_handler.py | 6 + litellm/proxy/utils.py | 17 ++- .../test_proxy_update_spend.py | 109 +++++++++++++++++- 3 files changed, 128 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 19bddee618b..f469587ab8e 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -199,6 +199,12 @@ class PrismaDBExceptionHandler: return True return False + @staticmethod + def is_prisma_error(e: Exception) -> bool: + import prisma + + return isinstance(e, _exception_types(prisma.errors.PrismaError)) + @staticmethod def is_deadlock_error(e: Exception) -> bool: """True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1f453d3b1ba..accf7b720fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6386,10 +6386,17 @@ class ProxyUpdateSpend: ) break except Exception as e: - if not PrismaDBExceptionHandler.is_database_transport_error(e): + if not _is_transient_spend_log_write_error(e): + if PrismaDBExceptionHandler.is_prisma_error(e): + await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + verbose_proxy_logger.warning( + "Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s", + len(logs_to_process), + str(e), + ) raise verbose_proxy_logger.warning( - "Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s", + "Spend tracking - transient DB error writing spend logs, retry %d/%d. logs_count=%d, error=%s", i + 1, n_retry_times, len(logs_to_process), @@ -6732,6 +6739,10 @@ async def _monitor_spend_logs_queue( MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256 +def _is_transient_spend_log_write_error(e: Exception) -> bool: + return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) + + async def _create_spend_logs_with_poison_isolation( repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]], @@ -6767,6 +6778,8 @@ async def _create_spend_logs_with_poison_isolation( raise if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): raise + if PrismaDBExceptionHandler.is_deadlock_error(e): + raise budget_left: Final = max(failure_budget - 1, 0) if len(rows) == 1: request_id: Final = rows[0].get("request_id") diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index 048fddb10d6..d671a4ffc1f 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -473,6 +473,110 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage( assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] +def _deadlock_error() -> Exception: + return _data_error( + 'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, ' + 'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })' + ) + + +@pytest.mark.asyncio +async def test_update_spend_logs_retries_deadlock_and_keeps_every_row( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """A 40P01 deadlock aborts the whole insert, so the same rows succeed on replay. + Before the fix the deadlock surfaced as a plain ``DataError`` and went through + poison-row isolation, which bisected the batch and dropped every row the + deadlock happened to hit as if Postgres had rejected it. + """ + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + create_many = AsyncMock(side_effect=[_deadlock_error(), _deadlock_error(), None]) + mock_prisma_client.db.litellm_spendlogs.create_many = create_many + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [] + + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + attempts = tuple( + tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list + ) + assert attempts == (("a", "b"), ("a", "b"), ("a", "b")) + assert mock_prisma_client.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_update_spend_logs_requeues_batch_once_deadlock_retries_exhaust( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """If every retry deadlocks, the batch goes back to the head of the queue for + the next flush instead of being dropped. + """ + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_deadlock_error()) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")] + + with pytest.raises(type(_deadlock_error())): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=1, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 2 + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_update_spend_logs_requeues_batch_on_non_transport_db_error( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + """A DB error that is neither transport nor deadlock (here P2021, the table is + gone mid-migration) is not retried in place, but the dequeued batch must not + be lost either: it goes back to the head of the queue so it lands once the + DB is healthy again. + """ + from prisma.errors import TableNotFoundError + + err = TableNotFoundError( + {"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}} + ) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")] + + with pytest.raises(TableNotFoundError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")], + ) + + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"] + + @pytest.mark.asyncio async def test_requeue_after_outage_drops_oldest_logs_past_the_byte_budget( mock_prisma_client: Any, make_spend_log_row: Any @@ -549,8 +653,9 @@ async def test_flush_returns_the_bytes_it_took_off_the_queue(mock_prisma_client: async def test_update_spend_logs_does_not_requeue_non_transport_failures( mock_prisma_client: Any, make_spend_log_row: Any ) -> None: - """Only transport failures are worth replaying. A rejection the DB will keep - rejecting must not be requeued, or it would wedge the queue forever. + """Only DB failures are worth replaying. A row the proxy itself cannot + serialize would fail the same way on every flush, so requeueing it would + wedge the head of the queue forever. """ mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=ValueError("bad payload")) proxy_logging = MagicMock() From 86038318ce408b4f63767a2fb753357ae7ed3cdf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 11:43:52 -0700 Subject: [PATCH 46/61] feat(ui): deep link guardrail detail with ?guardrail= on guardrails pages The Guardrails and Guardrails Monitor pages kept the selected guardrail in local React state, so the detail view could not be shared, reloaded, or reached with the browser back button. Both pages now read and write the selection through the nuqs `guardrail` query param, matching how the keys, teams, orgs, projects, users, models and logs pages deep link their detail views. Opening a guardrail pushes a history entry and closing it replaces the entry so back returns to the page the user came from --- .../GuardrailsMonitorView.test.tsx | 116 ++++++++++++++---- .../_components/GuardrailsMonitorView.tsx | 16 +-- .../_components/GuardrailsPanel.test.tsx | 77 +++++++++--- .../_components/GuardrailsPanel.tsx | 18 ++- 4 files changed, 179 insertions(+), 48 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx index 9f27daab6b3..4e0590df72d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.test.tsx @@ -1,36 +1,61 @@ -import { render, screen, waitFor } from "@testing-library/react"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { describe, expect, it, vi } from "vitest"; +import { type UrlUpdateEvent } from "nuqs/adapters/testing"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; import GuardrailsMonitorView from "./GuardrailsMonitorView"; import * as networking from "@/components/networking"; +import { renderWithProviders, screen, testQueryClient, waitFor } from "@/../tests/test-utils"; vi.mock("@/components/networking", () => ({ getGuardrailsUsageOverview: vi.fn(), + getGuardrailsUsageDetail: vi.fn(), + getGuardrailsUsageLogs: vi.fn(), formatDate: vi.fn((d: Date) => d.toISOString().slice(0, 10)), })); -const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview); +vi.mock("@/components/GuardrailsMonitor/LogViewer", () => ({ + LogViewer: ({ guardrailName }: { guardrailName: string }) =>
{guardrailName}
, +})); -function wrapper({ children }: { children: React.ReactNode }) { - const queryClient = new QueryClient({ - defaultOptions: { - queries: { retry: false }, - }, - }); - return {children}; -} +const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview); +const mockGetGuardrailsUsageDetail = vi.mocked(networking.getGuardrailsUsageDetail); +const mockGetGuardrailsUsageLogs = vi.mocked(networking.getGuardrailsUsageLogs); + +const emptyOverview = { rows: [], chart: [], totalRequests: 0, totalBlocked: 0, passRate: 100 }; + +const piiRow = { + id: "gr-pii", + name: "PII Guard", + type: "pii", + provider: "LiteLLM", + requestsEvaluated: 10, + failRate: 10, + status: "healthy" as const, + trend: "stable" as const, +}; + +const piiDetail = { + guardrail_name: "PII Guard", + description: "", + status: "healthy", + provider: "LiteLLM", + type: "pii", + requestsEvaluated: 10, + failRate: 10, + avgScore: 0.5, + avgLatency: 20, +}; describe("GuardrailsMonitorView", () => { - it("should render overview and fetch guardrails usage when accessToken is provided", async () => { - mockGetGuardrailsUsageOverview.mockResolvedValue({ - rows: [], - chart: [], - totalRequests: 0, - totalBlocked: 0, - passRate: 100, - }); + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + mockGetGuardrailsUsageOverview.mockResolvedValue(emptyOverview); + mockGetGuardrailsUsageDetail.mockResolvedValue(piiDetail); + mockGetGuardrailsUsageLogs.mockResolvedValue({ logs: [], total: 0 }); + }); - render(, { wrapper }); + it("should render overview and fetch guardrails usage when accessToken is provided", async () => { + renderWithProviders(); expect(await screen.findByRole("heading", { name: /Guardrails Monitor/i })).toBeInTheDocument(); await waitFor(() => { @@ -39,7 +64,54 @@ describe("GuardrailsMonitorView", () => { }); it("should render without crashing when accessToken is null", async () => { - render(, { wrapper }); + renderWithProviders(); expect(await screen.findByRole("heading", { name: /Guardrails Monitor/i })).toBeInTheDocument(); }); + + describe("guardrail detail deep link (?guardrail=)", () => { + it("should open the detail view directly from a ?guardrail= deep link", async () => { + renderWithProviders(, { searchParams: "?guardrail=gr-pii" }); + + expect(await screen.findByRole("heading", { name: "PII Guard" })).toBeInTheDocument(); + expect(mockGetGuardrailsUsageDetail).toHaveBeenCalledWith( + "test-token", + "gr-pii", + expect.any(String), + expect.any(String), + ); + expect(screen.queryByRole("heading", { name: /Guardrails Monitor/i })).not.toBeInTheDocument(); + }); + + it("should push ?guardrail= as a new history entry when a guardrail is selected", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + mockGetGuardrailsUsageOverview.mockResolvedValue({ ...emptyOverview, rows: [piiRow] }); + renderWithProviders(, { onUrlUpdate }); + + await user.click(await screen.findByRole("button", { name: "PII Guard" })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.get("guardrail")).toBe("gr-pii"); + expect(lastUpdate.options.history).toBe("push"); + expect(await screen.findByRole("heading", { name: "PII Guard" })).toBeInTheDocument(); + }); + + it("should clear ?guardrail= by replacing history when going back to the overview", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + renderWithProviders(, { + searchParams: "?guardrail=gr-pii", + onUrlUpdate, + }); + + await user.click(await screen.findByRole("button", { name: /back to overview/i })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.has("guardrail")).toBe(false); + expect(lastUpdate.options.history).toBe("replace"); + expect(await screen.findByRole("heading", { name: /Guardrails Monitor/i })).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx index a9acf3e6377..f90a46e19e4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsMonitorView.tsx @@ -1,12 +1,11 @@ import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; +import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useMemo, useState } from "react"; import { formatDate } from "@/components/networking"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { GuardrailDetail } from "./GuardrailDetail"; import { GuardrailsOverview } from "./GuardrailsOverview"; -type View = { type: "overview" } | { type: "detail"; guardrailId: string }; - interface GuardrailsMonitorViewProps { accessToken?: string | null; } @@ -16,7 +15,10 @@ const defaultStart = new Date(); defaultStart.setDate(defaultStart.getDate() - 7); export default function GuardrailsMonitorView({ accessToken = null }: GuardrailsMonitorViewProps) { - const [view, setView] = useState({ type: "overview" }); + const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( + "guardrail", + parseAsString.withOptions({ history: "push" }), + ); const initialFrom = useMemo(() => new Date(defaultStart), []); const initialTo = useMemo(() => new Date(defaultEnd), []); @@ -34,11 +36,11 @@ export default function GuardrailsMonitorView({ accessToken = null }: Guardrails }, []); const handleSelectGuardrail = (id: string) => { - setView({ type: "detail", guardrailId: id }); + void setSelectedGuardrailId(id); }; const handleBack = () => { - setView({ type: "overview" }); + void setSelectedGuardrailId(null, { history: "replace" }); }; const dateRangeControl = ( @@ -47,7 +49,7 @@ export default function GuardrailsMonitorView({ accessToken = null }: Guardrails return (
- {view.type === "overview" ? ( + {!selectedGuardrailId ? (
{dateRangeControl}
({ getGuardrailsList: vi.fn(), @@ -15,16 +16,21 @@ vi.mock("./add_guardrail_form", () => ({ vi.mock("./guardrail_table", () => ({ __esModule: true, - default: ({ guardrailsList, onDeleteClick }: any) => ( + default: ({ guardrailsList, onDeleteClick, onGuardrailClick }: any) => (
Mock Guardrail Table
{guardrailsList.length > 0 && ( - + <> + + + )}
), @@ -32,7 +38,12 @@ vi.mock("./guardrail_table", () => ({ vi.mock("./guardrail_info", () => ({ __esModule: true, - default: () =>
Mock Guardrail Info View
, + default: ({ guardrailId, onClose }: { guardrailId: string; onClose: () => void }) => ( +
+
Mock Guardrail Info View {guardrailId}
+ +
+ ), })); vi.mock("./GuardrailTestPlayground", async () => { @@ -112,7 +123,7 @@ describe("GuardrailsPanel", () => { }); it("should render the component", async () => { - render(); + renderWithProviders(); expect(screen.getByText("Guardrails")).toBeInTheDocument(); // Activate the Guardrails tab so its content (including the Add button) is rendered fireEvent.click(screen.getByText("Guardrails")); @@ -120,7 +131,7 @@ describe("GuardrailsPanel", () => { }); it("should delete the clicked guardrail after confirming in the modal", async () => { - render(); + renderWithProviders(); fireEvent.click(screen.getByText("Guardrails")); fireEvent.click(await screen.findByTestId("delete-button")); @@ -139,14 +150,14 @@ describe("GuardrailsPanel", () => { }); it("should mount every tab panel up front so panel state survives tab switches", async () => { - render(); + renderWithProviders(); expect(await screen.findByLabelText("playground draft")).toBeInTheDocument(); expect(screen.getByText("Mock Team Guardrails Tab")).toBeInTheDocument(); }); it("should keep test playground state when switching tabs away and back", async () => { - render(); + renderWithProviders(); fireEvent.click(screen.getByText("Test Playground")); @@ -161,7 +172,7 @@ describe("GuardrailsPanel", () => { }); it("should not delete anything when the modal is cancelled", async () => { - render(); + renderWithProviders(); fireEvent.click(screen.getByText("Guardrails")); fireEvent.click(await screen.findByTestId("delete-button")); @@ -171,4 +182,42 @@ describe("GuardrailsPanel", () => { expect(mockDeleteGuardrailCall).not.toHaveBeenCalled(); }); + + describe("guardrail detail deep link (?guardrail=)", () => { + it("should open the guardrail info view directly from a ?guardrail= deep link", async () => { + renderWithProviders(, { searchParams: "?guardrail=test-guardrail-1" }); + + expect(await screen.findByTestId("guardrail-info-view")).toHaveTextContent("test-guardrail-1"); + expect(screen.queryByText("Mock Guardrail Table")).not.toBeInTheDocument(); + }); + + it("should push ?guardrail= as a new history entry when a guardrail row is clicked", async () => { + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + renderWithProviders(, { onUrlUpdate }); + + fireEvent.click(await screen.findByTestId("open-button")); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.get("guardrail")).toBe("test-guardrail-1"); + expect(lastUpdate.options.history).toBe("push"); + expect(await screen.findByTestId("guardrail-info-view")).toHaveTextContent("test-guardrail-1"); + }); + + it("should clear ?guardrail= by replacing history when the info view is closed", async () => { + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + renderWithProviders(, { + searchParams: "?guardrail=test-guardrail-1", + onUrlUpdate, + }); + + fireEvent.click(await screen.findByRole("button", { name: "Close Guardrail Info" })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const lastUpdate = onUrlUpdate.mock.calls.at(-1)![0]; + expect(lastUpdate.searchParams.has("guardrail")).toBe(false); + expect(lastUpdate.options.history).toBe("replace"); + expect(await screen.findByText("Mock Guardrail Table")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index 7e59abf8e3d..901e39004f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -1,3 +1,4 @@ +import { parseAsString, useQueryState } from "nuqs"; import React, { useState, useEffect } from "react"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { ChevronDown, Code, Plus } from "lucide-react"; @@ -40,7 +41,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [selectedGuardrailId, setSelectedGuardrailId] = useState(null); + const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( + "guardrail", + parseAsString.withOptions({ history: "push" }), + ); const isAdmin = userRole ? isAdminRole(userRole) : false; const fetchGuardrails = async () => { @@ -63,16 +67,20 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole fetchGuardrails(); }, [accessToken]); + const closeGuardrailDetail = () => { + void setSelectedGuardrailId(null, { history: "replace" }); + }; + const handleAddGuardrail = () => { if (selectedGuardrailId) { - setSelectedGuardrailId(null); + closeGuardrailDetail(); } setIsAddModalVisible(true); }; const handleAddCustomCodeGuardrail = () => { if (selectedGuardrailId) { - setSelectedGuardrailId(null); + closeGuardrailDetail(); } setIsCustomCodeModalVisible(true); }; @@ -175,7 +183,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole {selectedGuardrailId ? ( setSelectedGuardrailId(null)} + onClose={closeGuardrailDetail} accessToken={accessToken} isAdmin={isAdmin} /> @@ -184,7 +192,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole guardrailsList={guardrailsList} isLoading={isLoading} onDeleteClick={handleDeleteClick} - onGuardrailClick={(id) => setSelectedGuardrailId(id)} + onGuardrailClick={(id) => void setSelectedGuardrailId(id)} /> )} From e3b4a82ff9991369f0a79e34f44a5da732506b0f Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 11:44:34 -0700 Subject: [PATCH 47/61] Merge pull request #39926 from BerriAI/litellm_lit6981_none_url_auth fix(mcp): reject URL credentials for none auth --- .../_experimental/mcp_server/exceptions.py | 14 +++++++ .../outbound_credentials/adapter.py | 3 ++ .../outbound_credentials/resolver.py | 11 ++++- .../mcp_server/outbound_credentials/types.py | 12 ++++++ .../mcp_server/rest_endpoints.py | 3 ++ .../outbound_credentials/test_adapter.py | 12 ++++++ .../outbound_credentials/test_resolver.py | 29 +++++++++++++ .../outbound_credentials/test_types.py | 9 ++++ .../mcp_server/test_mcp_server_manager.py | 42 +++++++++++++++++++ .../mcp_server/test_rest_endpoints.py | 30 +++++++++++++ 10 files changed, 164 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/exceptions.py b/litellm/proxy/_experimental/mcp_server/exceptions.py index a1b3b167a4a..c818f6b05bd 100644 --- a/litellm/proxy/_experimental/mcp_server/exceptions.py +++ b/litellm/proxy/_experimental/mcp_server/exceptions.py @@ -5,6 +5,20 @@ from typing import Final from fastapi import HTTPException +class MCPServerURLCredentialsError(HTTPException): + """A fixed, sanitized URL-credential migration error safe for operator previews.""" + + def __init__(self) -> None: + super().__init__( + status_code=500, + detail=( + "misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; " + "remove them from the URL and configure Basic Auth with auth_type: basic and " + "auth_value: username:password" + ), + ) + + class MCPUpstreamAuthError(Exception): """Raised when an upstream MCP server returns an authentication failure (typically HTTP 401) and the gateway should surface it transparently to diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 6a95a93a2a8..ea2318bd6f1 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -19,6 +19,7 @@ from pydantic import SecretStr from typing_extensions import assert_never from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials +from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, @@ -293,6 +294,8 @@ def raise_public(error: CredError) -> NoReturn: ) case "misconfigured": raise HTTPException(status_code=500, detail=error.summary) + case "url_credentials_not_allowed": + raise MCPServerURLCredentialsError() case "upstream_unavailable": raise HTTPException(status_code=503, detail=error.summary) case "unsupported_mode": diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 3af7b51f432..404baa14350 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -134,7 +134,7 @@ class UpstreamCredentialProvider: async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: case NoneConfig(): - return Ok(NoOpAuth()) + return self._none(server) case ApiKeyConfig() as config: return self._api_key(config) case PassthroughConfig(): @@ -151,6 +151,15 @@ class UpstreamCredentialProvider: return _not_implemented(AuthSpecKind.aws_sigv4) assert_never(server.config) + def _none(self, server: ServerSpec) -> Result[httpx.Auth, CredError]: + try: + resource: Final = httpx.URL(server.resource) + except httpx.InvalidURL: + return Ok(NoOpAuth()) + if resource.userinfo: + return Error(CredError.of_url_credentials_not_allowed()) + return Ok(NoOpAuth()) + async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool: """Whether a usable per-user token exists for this server (the preemptive 401's check). diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py index 67aad3e443e..632dc57dcf6 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py @@ -95,6 +95,7 @@ class CredError: tag: Literal[ "unauthorized", "misconfigured", + "url_credentials_not_allowed", "upstream_unavailable", "unsupported_mode", "precondition_required", @@ -103,6 +104,7 @@ class CredError: unauthorized: Unauthorized = case() # no usable credential for this (subject, server) -> 401 challenge misconfigured: str = case() # the declared mode is missing required config -> 5xx (operator) + url_credentials_not_allowed: None = case() upstream_unavailable: str = case() # the IdP / token endpoint could not be reached -> 503 unsupported_mode: str = case() # a raw mode string did not parse into AuthSpecKind (boundary) precondition_required: str = case() # a required per-user value (e.g. an env var) has not been provided -> 412 @@ -129,6 +131,10 @@ class CredError: def of_misconfigured(detail: str) -> CredError: return CredError(misconfigured=detail) + @staticmethod + def of_url_credentials_not_allowed() -> CredError: + return CredError(url_credentials_not_allowed=None) + @staticmethod def of_upstream_unavailable(detail: str) -> CredError: return CredError(upstream_unavailable=detail) @@ -154,6 +160,12 @@ class CredError: return f"unauthorized: {self.unauthorized.detail}" case "misconfigured": return f"misconfigured: {self.misconfigured}" + case "url_credentials_not_allowed": + return ( + "misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; " + "remove them from the URL and configure Basic Auth with auth_type: basic and " + "auth_value: username:password" + ) case "upstream_unavailable": return f"upstream unavailable: {self.upstream_unavailable}" case "unsupported_mode": diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index b3469da9071..5fbfad54a39 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -18,6 +18,7 @@ from litellm.exceptions import ( ) from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, + MCPServerURLCredentialsError, MCPUpstreamAuthError, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( @@ -75,6 +76,8 @@ _MCP_GUARDRAIL_REJECTIONS: Final = ( def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str: + if isinstance(exc, MCPServerURLCredentialsError): + return str(exc.detail) if isinstance(exc, TimeoutError): return ( f"Failed to connect to MCP server: no response from {url or 'the server'} " diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index c6f3b9cb1f4..d67d0df4d0e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -13,6 +13,7 @@ from fastapi import HTTPException from pydantic import ValidationError from litellm.experimental_mcp_client.client import MCPClient +from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( oauth_protected_resource_path, raise_public, @@ -442,6 +443,17 @@ def test_raise_public_maps_each_error_to_its_status(error, status): assert exc_info.value.status_code == status +def test_raise_public_marks_only_url_credentials_error_as_safe_for_preview(): + with pytest.raises(HTTPException) as generic_exc_info: + raise_public(CredError.of_misconfigured("private operator detail")) + assert not isinstance(generic_exc_info.value, MCPServerURLCredentialsError) + + error = CredError.of_url_credentials_not_allowed() + with pytest.raises(MCPServerURLCredentialsError) as url_exc_info: + raise_public(error) + assert url_exc_info.value.detail == error.summary + + def test_raise_public_emits_unauthorized_challenge(): body = {"error": "byok_auth_required", "server_id": "s1"} error = CredError.of_unauthorized("needs key", www_authenticate='Bearer resource_metadata="/x"', body=body) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 9d63e8c2c1c..5e2f2cf97d7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -116,6 +116,35 @@ async def test_none_mode_yields_a_no_op_auth(): assert isinstance(result.ok, NoOpAuth) +@pytest.mark.asyncio +async def test_none_mode_rejects_url_userinfo(): + spec = ServerSpec( + server_id="s", + resource="https://lit-user:s3cr3t@upstream.example.com/mcp", + config=NoneConfig(), + ) + + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec) + + assert isinstance(result, Error) + assert result.error.tag == "url_credentials_not_allowed" + assert "Basic Auth" in result.error.summary + assert "auth_type: basic" in result.error.summary + assert "auth_value: username:password" in result.error.summary + assert "lit-user" not in result.error.summary + assert "s3cr3t" not in result.error.summary + + +@pytest.mark.asyncio +async def test_none_mode_does_not_validate_non_credential_resource(): + spec = ServerSpec(server_id="s", resource="https://[::1", config=NoneConfig()) + + result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec) + + assert isinstance(result, Ok) + assert isinstance(result.ok, NoOpAuth) + + @pytest.mark.asyncio async def test_api_key_shared_emits_the_configured_header(): config = ApiKeyConfig( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py index d4b51b08e06..bacbb5c1236 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py @@ -75,6 +75,15 @@ def test_crederror_factory_sets_the_matching_tag(factory, expected_tag): assert "detail text" in err.summary +def test_url_credentials_error_has_a_fixed_actionable_summary(): + err = CredError.of_url_credentials_not_allowed() + + assert err.tag == "url_credentials_not_allowed" + assert "Basic Auth" in err.summary + assert "auth_type: basic" in err.summary + assert "auth_value: username:password" in err.summary + + def test_apikeyconfig_requires_a_key_source(): with pytest.raises(ValidationError): ApiKeyConfig() # type: ignore[call-arg] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9745e508703..34dc067e7a3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -8966,6 +8966,24 @@ class TestCreateMcpClientV2Graft: assert isinstance(client._resolved_auth, NoOpAuth) assert client._mcp_auth_value is None + @pytest.mark.parametrize("auth_type", [None, MCPAuth.none]) + async def test_none_mode_rejects_url_userinfo(self, auth_type): + with pytest.raises(HTTPException) as exc_info: + await MCPServerManager()._create_mcp_client( + self._http_server( + auth_type=auth_type, + url="https://lit-user:s3cr3t@upstream.example.com/mcp", + ) + ) + + detail = str(exc_info.value.detail) + assert exc_info.value.status_code == 500 + assert "Basic Auth" in detail + assert "auth_type: basic" in detail + assert "auth_value: username:password" in detail + assert "lit-user" not in detail + assert "s3cr3t" not in detail + @pytest.mark.parametrize( "auth_type, token, expected_name, expected_value", [ @@ -11369,6 +11387,30 @@ class TestResolveOpenapiToolAuth: assert "Authorization" not in (forwarded or {}) + @pytest.mark.asyncio + async def test_none_mode_without_url_keeps_spec_path_server_unauthenticated(self): + server = MCPServer( + server_id="openapi-only", + name="report_api", + server_name="report_api", + url=None, + transport=MCPTransport.http, + auth_type=MCPAuth.none, + spec_path="https://api.example.com/openapi.json", + ) + + resolved, forwarded = await MCPServerManager().resolve_openapi_upstream_auth( + mcp_server=server, + oauth2_headers=None, + raw_headers=None, + mcp_auth_header=None, + user_api_key_auth=None, + forwarded_headers={"X-Trace": "trace-id"}, + ) + + assert resolved is None + assert forwarded == {"X-Trace": "trace-id"} + class TestOpenApiHandlerRelaysUpstreamAuth: """`_call_openapi_tool_handler` must not flatten a re-auth signal into a generic message. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index f2c8f8c80c5..c007d22117f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -177,6 +177,36 @@ class TestExecuteWithMcpClient: assert "https://api.example.com/mcp/" in message assert "30s" in message + def test_connection_error_message_hides_arbitrary_http_exception_detail(self): + message = rest_endpoints._connection_error_message( + HTTPException(status_code=500, detail="secret upstream detail"), + "https://api.example.com/mcp/", + 30.0, + ) + + assert "secret upstream detail" not in message + + @pytest.mark.asyncio + async def test_none_mode_url_credentials_returns_actionable_redacted_error(self): + async def unreached_operation(client): + raise AssertionError("operation must not run for an invalid server configuration") + + payload = NewMCPServerRequest( + server_name="example", + url="https://lit-user:s3cr3t@upstream.example.com/mcp", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client(payload, unreached_operation) + + message = str(result["message"]) + assert result["error"] is True + assert "Basic Auth" in message + assert "auth_type: basic" in message + assert "auth_value: username:password" in message + assert "lit-user" not in message + assert "s3cr3t" not in message + @pytest.mark.asyncio async def test_forwards_static_headers(self, monkeypatch): """Ensure static_headers are forwarded to the MCP client during test calls. From 9ac893df1562f35a56fbac17399481f1fb32f857 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 11:47:15 -0700 Subject: [PATCH 48/61] style(e2e): wrap the openai passthrough content assertion under 120 columns --- tests/e2e/llm_translation/test_passthrough_e2e.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index e50e83eaf77..447fe7d30d9 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -361,8 +361,14 @@ class TestOpenAIProviderPrefixChat: completion = ChatResponse.model_validate_json(result.body) assert completion.id, f"/openai/v1/chat/completions relayed no completion id: {result.body[:300]}" - content = completion.choices[0].message.content if completion.choices and completion.choices[0].message else None - assert content and content.strip(), f"/openai/v1/chat/completions relayed an empty completion: {result.body[:300]}" + content = ( + completion.choices[0].message.content + if completion.choices and completion.choices[0].message + else None + ) + assert content and content.strip(), ( + f"/openai/v1/chat/completions relayed an empty completion: {result.body[:300]}" + ) assert completion.usage is not None, f"the completion carried no usage to price from: {completion}" row = _fetch_cost_breakdown(client, completion.id) From a0058ed15759febc3acb96ab27c823671a232715 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 5 Sep 2026 11:47:36 -0700 Subject: [PATCH 49/61] fix(hide-secrets): stop redacting benign identifiers (#39879) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(hide-secrets): stop redacting benign identifiers and make redaction deterministic The OpenAI key detector matched `sk-` anywhere inside a word, so `` became ``, and the Base64 entropy limit of 3.0 flagged ordinary quoted identifiers such as `"application/json"` and model ids. Redaction also iterated a hash-seeded set, so the same request produced different bytes on different workers and broke prompt caching. - require a standalone `sk-`/`sk_` token with a digit (still catches sk-proj-/sk-ant-) - raise Base64HighEntropyString limit from 3.0 to the detect-secrets default 4.5 - redact overlapping matches longest-first in a stable order Resolves LIT-7049 * fix(hide-secrets): treat separators as key boundaries and defer sk_live_ to the stripe detector The standalone-token boundary also rejected keys glued to a preceding `_`, `-` or percent-encoded delimiter (`openai_sk-…`, `key-sk-…`, `Bearer%20sk-…`), which the old pattern redacted, and `sk_live_…` was counted by both the OpenAI and the Stripe detector. * fix(hide-secrets): keep the openai key scan linear on repeated sk separators The digit requirement was a lookahead, so every `sk` inside a long `[a-zA-Z0-9_-]` run re-scanned the rest of that run looking for a digit. 100 KB of `-sk-` took over 5s in the worker's event loop and the proxy closed the connection without a response. The check now runs once per match in `analyze_string` instead. * chore(hide-secrets): remove redundant performance test comment * fix(hide-secrets): consume complete openai key tokens * chore(hide-secrets): remove redundant fixture comment * chore(hide-secrets): remove redundant test docstrings * fix(hide-secrets): redact whole stripe live keys * style(hide-secrets): wrap secret sorting key --- .../enterprise_callbacks/secret_detection.py | 27 ++--- .../secrets_plugins/openai_api_key.py | 15 ++- .../test_secret_detection.py | 101 ++++++++++++++++-- 3 files changed, 124 insertions(+), 19 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 1fddc527ec8..bfbfd7bfb15 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -433,9 +433,9 @@ _default_detect_secrets_config = { "name": "ZendeskSecretKeyDetector", "path": _custom_plugins_path + "/zendesk_secret_key.py", }, - {"name": "Base64HighEntropyString", "limit": 3.0}, + {"name": "Base64HighEntropyString", "limit": 4.5}, {"name": "HexHighEntropyString", "limit": 3.0}, - ] + ], } @@ -466,16 +466,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): os.remove(temp_file.name) - detected_secrets = [] - for file in secrets.files: - for found_secret in secrets[file]: - if found_secret.secret_value is None: - continue - detected_secrets.append( - {"type": found_secret.type, "value": found_secret.secret_value} - ) - - return detected_secrets + return [ + {"type": found_secret.type, "value": found_secret.secret_value} + for file in sorted(secrets.files) + for found_secret in sorted( + secrets[file], + key=lambda secret: ( + -len(secret.secret_value or ""), + secret.type, + secret.secret_value or "", + ), + ) + if found_secret.secret_value is not None + ] def redact_text(self, text: str, source: str = "message") -> str: """Replace every detected secret in ``text`` with ``[REDACTED]`` and diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py b/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py index c5d20f75909..32652703326 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py @@ -3,6 +3,7 @@ This plugin searches for OpenAI API Keys. """ import re +from collections.abc import Generator from detect_secrets.plugins.base import RegexBasedDetector @@ -16,4 +17,16 @@ class OpenAIApiKeyDetector(RegexBasedDetector): @property def denylist(self) -> list[re.Pattern]: - return [re.compile(r"""(sk-[a-zA-Z0-9]{5,})""")] + return [ + re.compile( + r"((?:(? Generator[str, None, None]: + # the digit check lives outside the regex: a lookahead re-scans the token + # from every `sk` inside it, which is quadratic on `-sk-sk-sk-...` input + yield from (match for match in super().analyze_string(string) if re.search(r"[0-9]", match)) diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index dc1cbb9983e..f46df5baadf 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -10,6 +10,8 @@ Covers the three defects from the ticket: handling live only on the native path). """ +import time + import pytest from litellm_enterprise.enterprise_callbacks.secret_detection import ( @@ -19,12 +21,16 @@ from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth AWS_KEY = "AKIAIOSFODNN7EXAMPLE" +OPENAI_KEY = "sk-test-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH" +SHORT_OPENAI_KEY = "sk-12345" +UNICODE_DIGIT_SUFFIX = "sk-notification٣" +STRIPE_LIVE_KEY = f"sk_live_{'1234567890' * 3}" +URL_ENCODED_KEY = "Bearer%20sk-Ab3dEf6Gh7Ij8Kl9Mn0Pq2Rs3Tu4Vw5X" +AWS_KEYS = [f"AKIAIOSFODNN7EXAMPL{suffix}" for suffix in "FEDCBA"] def _guardrail() -> _ENTERPRISE_SecretDetection: - return _ENTERPRISE_SecretDetection( - guardrail_name="hide-secrets", event_hook="pre_call", default_on=True - ) + return _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets", event_hook="pre_call", default_on=True) def _recorded(request_data: dict) -> dict: @@ -33,6 +39,91 @@ def _recorded(request_data: dict) -> dict: return entries[0] +def test_scan_message_preserves_benign_identifiers_and_xml_tags(): + guardrail = _guardrail() + content = " model: claude-sonnet-4-5-20250929 " + + assert guardrail.scan_message_for_secrets(content) == [] + assert guardrail.redact_text(content) == content + assert guardrail.redact_text("result = compute(x) ") == ( + "result = compute(x) " + ) + + +def test_scan_message_preserves_quoted_benign_identifiers(): + guardrail = _guardrail() + content = '{"content-type": "application/json", "model": "claude-sonnet-4-5-20250929"}' + + assert guardrail.scan_message_for_secrets(content) == [] + assert guardrail.redact_text(content) == content + + +def test_scan_message_redacts_every_openai_key_occurrence(): + guardrail = _guardrail() + content = f"first {OPENAI_KEY}, second {OPENAI_KEY}" + + assert guardrail.redact_text(content) == "first [REDACTED], second [REDACTED]" + + +def test_scan_message_redacts_short_numeric_openai_like_values(): + guardrail = _guardrail() + + assert guardrail.redact_text(f"value {SHORT_OPENAI_KEY}") == "value [REDACTED]" + + +def test_scan_message_requires_ascii_digits_for_openai_like_values(): + guardrail = _guardrail() + + assert guardrail.scan_message_for_secrets(UNICODE_DIGIT_SUFFIX) == [] + assert guardrail.redact_text(UNICODE_DIGIT_SUFFIX) == UNICODE_DIGIT_SUFFIX + + +def test_scan_message_redacts_openai_key_after_separator(): + guardrail = _guardrail() + + assert guardrail.redact_text(f"openai_{OPENAI_KEY} key-{OPENAI_KEY}") == ( + "openai_[REDACTED] key-[REDACTED]" + ) + assert guardrail.redact_text(URL_ENCODED_KEY) == "Bearer%20[REDACTED]" + + +def test_scan_message_does_not_stop_openai_key_at_token_characters(): + guardrail = _guardrail() + + assert guardrail.redact_text("key sk-proj-abcde12345/extra") == "key [REDACTED]/extra" + + +def test_scan_message_stays_linear_on_repeated_sk_separators(): + guardrail = _guardrail() + content = "-sk-" * 25_000 + + started = time.perf_counter() + assert guardrail.scan_message_for_secrets(content) == [] + assert time.perf_counter() - started < 2.0 + + +def test_scan_message_redacts_whole_stripe_live_key(): + guardrail = _guardrail() + + assert guardrail.redact_text(f"stripe {STRIPE_LIVE_KEY} end") == "stripe [REDACTED] end" + + +def test_scan_message_returns_matches_in_stable_order(): + guardrail = _guardrail() + detected = guardrail.scan_message_for_secrets(" ".join(AWS_KEYS)) + + assert [secret["value"] for secret in detected] == sorted(AWS_KEYS) + + +def test_scan_message_replaces_longest_overlapping_match_first(): + guardrail = _guardrail() + content = f'token = "{OPENAI_KEY}/extra"' + + detected = guardrail.scan_message_for_secrets(content) + assert [secret["value"] for secret in detected] == [f"{OPENAI_KEY}/extra", OPENAI_KEY] + assert guardrail.redact_text(content) == 'token = "[REDACTED]"' + + @pytest.mark.asyncio async def test_apply_guardrail_redacts_secrets(): """Playground path: the returned texts must carry [REDACTED], not the secret.""" @@ -199,9 +290,7 @@ async def test_apply_guardrail_without_texts_records_nothing(): "messages": [ { "role": "user", - "content": [ - {"type": "image_url", "image_url": {"url": "https://x/y.png"}} - ], + "content": [{"type": "image_url", "image_url": {"url": "https://x/y.png"}}], } ], "metadata": {}, From 3c0900b7c5d26aec0b6ed508083dcee20d6501e2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:51:15 -0700 Subject: [PATCH 50/61] perf(logging): scan large base64 payloads for log truncation off the event loop (#39890) * perf(logging): scan large base64 payloads for log truncation off the event loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * perf(logging): make base64 offload threshold a plain constant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 21 ++++-- litellm/litellm_core_utils/logging_utils.py | 40 ++++++++++- .../test_litellm_logging.py | 50 +++++++++++++ .../litellm_core_utils/test_logging_utils.py | 71 +++++++++++++++++++ 5 files changed, 177 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e4fb2a00297..ce744e9c58a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -89,6 +89,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = ( # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)) +BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: Final = 256 * 1024 REDACTED_BY_LITELLM: Final = "redacted-by-litellm" # in-memory stand-in handed to provider converters for redacted arguments; never stored REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8f1b2ce1cc2..01b823e51ab 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -78,7 +78,10 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( InteractionsUsageObjectTransformation, ) -from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages +from litellm.litellm_core_utils.logging_utils import ( + truncate_base64_in_messages, + truncate_base64_in_messages_async, +) from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.redact_messages import ( redact_message_input_output_from_custom_logger, @@ -538,6 +541,7 @@ class Logging(LiteLLMLoggingBaseClass): self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( self.initialize_standard_built_in_tools_params(kwargs) ) + self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape ## TIME TO FIRST TOKEN LOGGING ## self.completion_start_time: datetime.datetime | None = None self._llm_caching_handler: LLMCachingHandler | None = None @@ -2933,6 +2937,11 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage + self.truncated_messages_for_logging = await truncate_base64_in_messages_async( + StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=self.model_call_details, messages=self.model_call_details.get("messages") + ) + ) start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, end_time=end_time, @@ -6202,9 +6211,13 @@ def get_standard_logging_object_payload( model_id=_model_id, requester_ip_address=clean_metadata.get("requester_ip_address", None), user_agent=clean_metadata.get("user_agent", None), - messages=truncate_base64_in_messages( - StandardLoggingPayloadSetup.append_system_prompt_messages( - kwargs=kwargs, messages=kwargs.get("messages") + messages=( + logging_obj.truncated_messages_for_logging + if logging_obj.truncated_messages_for_logging is not None + else truncate_base64_in_messages( + StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=kwargs, messages=kwargs.get("messages") + ) ) ), response=final_response_obj, diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index f3b1b29a9ad..44daef42e14 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -3,12 +3,15 @@ import functools import inspect import re import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger -from litellm.constants import MAX_BASE64_LENGTH_FOR_LOGGING +from litellm.constants import ( + BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS, + MAX_BASE64_LENGTH_FOR_LOGGING, +) from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -141,6 +144,39 @@ def truncate_base64_in_messages( return messages +_StringTree = str | Sequence["_StringTree"] | Mapping[str, "_StringTree"] | None + + +def _iter_string_leaves(value: _StringTree) -> Iterator[str]: + stack: Final[list[_StringTree]] = [value] # mutable-ok: explicit stack, recursive functions are banned in litellm/ + while stack: + match stack.pop(): + case str() as text: + yield text + case Mapping() as mapping: + stack.extend(mapping.values()) + case Sequence() as items: + stack.extend(items) + case None: + pass + + +async def truncate_base64_in_messages_async( + messages: str | list | dict | None, # mutable-ok: same contract as truncate_base64_in_messages +) -> str | list | dict | None: # mutable-ok: same contract as truncate_base64_in_messages + """ + Same result as truncate_base64_in_messages, but payloads whose string content + reaches BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are scanned in a worker + thread so the regex pass over multi-MB base64 images does not block the event loop. + """ + if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0: + return messages + total_chars: Final = sum(len(leaf) for leaf in _iter_string_leaves(messages)) + if total_chars < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: + return truncate_base64_in_messages(messages) + return await asyncio.to_thread(truncate_base64_in_messages, messages) + + # Global service logger instance to avoid recreating it _service_logger = None diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index af75691eb10..31fb4fb55c5 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1114,6 +1114,56 @@ async def test_logging_non_streaming_request(): litellm.callbacks = original_callbacks +@pytest.mark.asyncio +async def test_async_success_handler_truncates_large_base64_off_the_event_loop(monkeypatch): + """The standard logging payload's base64 scan of a large multimodal request must not run on the loop thread.""" + import threading + + from litellm.litellm_core_utils import logging_utils + + loop_thread = threading.get_ident() + scan_threads: list[int] = [] + original_scan = logging_utils._truncate_base64_in_string + + def recording_scan(value: str) -> str: + scan_threads.append(threading.get_ident()) + return original_scan(value) + + monkeypatch.setattr(logging_utils, "_truncate_base64_in_string", recording_scan) + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + + logged = asyncio.Event() + captured: dict = {} + + class CaptureLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + captured["standard_logging_object"] = kwargs["standard_logging_object"] + logged.set() + + monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()]) + payload = "L" * 20_000 + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{payload}"}}, + ], + } + ], + mock_response="ok", + ) + await asyncio.wait_for(logged.wait(), timeout=10) + + logged_url = captured["standard_logging_object"]["messages"][0]["content"][1]["image_url"]["url"] + assert "base64_data truncated" in logged_url + assert payload not in logged_url + assert scan_threads + assert loop_thread not in scan_threads + + @pytest.mark.parametrize( "async_flag", [ diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/test_litellm/litellm_core_utils/test_logging_utils.py index b0dad0bf228..f9913f1935d 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_utils.py @@ -2,12 +2,16 @@ Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers. """ +import threading + import pytest +from litellm.litellm_core_utils import logging_utils from litellm.litellm_core_utils.logging_utils import ( _format_base64_size, _truncate_base64_in_string, truncate_base64_in_messages, + truncate_base64_in_messages_async, ) # --------------------------------------------------------------------------- @@ -157,3 +161,70 @@ class TestTruncateBase64InMessages: result[0]["content"][0]["image_url"]["url"] == f"data:image/png;base64,{short}" ) + + +# --------------------------------------------------------------------------- +# truncate_base64_in_messages_async +# --------------------------------------------------------------------------- + + +def _image_messages(payload: str) -> list: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{payload}"}}, + ], + } + ] + + +@pytest.fixture +def scan_threads(monkeypatch): + """Record the thread that runs every base64 regex scan.""" + threads: list[int] = [] + original = logging_utils._truncate_base64_in_string + + def recording_scan(value: str) -> str: + threads.append(threading.get_ident()) + return original(value) + + monkeypatch.setattr(logging_utils, "_truncate_base64_in_string", recording_scan) + return threads + + +class TestTruncateBase64InMessagesAsync: + @pytest.mark.asyncio + async def test_large_payload_is_scanned_off_the_event_loop(self, monkeypatch, scan_threads): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + payload = "I" * 20_000 + messages = _image_messages(payload) + + result = await truncate_base64_in_messages_async(messages) + offload_threads = tuple(scan_threads) + + assert result == truncate_base64_in_messages(messages) + assert payload not in result[0]["content"][1]["image_url"]["url"] + assert payload in messages[0]["content"][1]["image_url"]["url"] + assert offload_threads + assert threading.get_ident() not in offload_threads + + @pytest.mark.asyncio + async def test_small_payload_stays_on_the_calling_thread(self, monkeypatch, scan_threads): + monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000) + messages = _image_messages("J" * 200) + + result = await truncate_base64_in_messages_async(messages) + + assert result == truncate_base64_in_messages(messages) + assert scan_threads + assert set(scan_threads) == {threading.get_ident()} + + @pytest.mark.asyncio + async def test_none_and_disabled_truncation_short_circuit(self, monkeypatch, scan_threads): + assert await truncate_base64_in_messages_async(None) is None + monkeypatch.setattr(logging_utils, "MAX_BASE64_LENGTH_FOR_LOGGING", 0) + messages = _image_messages("K" * 20_000) + assert await truncate_base64_in_messages_async(messages) is messages + assert scan_threads == [] From a670a4621e9029b054de86ad28c0fe939d1cfc52 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 11:53:04 -0700 Subject: [PATCH 51/61] fix(proxy): make the invalid-model 403 path cheap under a burst of rejections (#39892) * fix(proxy): make the invalid-model 403 path cheap under a burst of rejections Keep the wildcard pattern registry in specificity order at registration time so route() no longer re-sorts every pattern per lookup, and reuse the standardized failure payload across the async and threaded sync failure handlers regardless of what a callback did to log_event_type. Rejections are still logged and observable. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(router): wrap the filtered pattern tuple the way ruff format wants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router,logging): assert registry order and callback awaits instead of patching a class Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): inject the pattern sorter so the lookup test observes that route() never sorts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 3 +- .../router_utils/pattern_match_deployments.py | 17 ++++++---- .../test_litellm_logging.py | 28 ++++++++++++++++ .../test_pattern_match_deployments.py | 32 ++++++++++++++++++- 4 files changed, 70 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 01b823e51ab..22e4dbf3a44 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3234,8 +3234,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details = {} if ( - self.model_call_details.get("log_event_type") == "failed_api_call" - and self.model_call_details.get("exception") is exception + self.model_call_details.get("exception") is exception and self.model_call_details.get("standard_logging_object") is not None ): return start_time, self.model_call_details["end_time"] diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 0775e0a4039..d5234e27ec6 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -56,8 +56,9 @@ class PatternMatchRouter: This class will store a mapping for regex pattern: List[Deployments] """ - def __init__(self): + def __init__(self, pattern_utils: type[PatternUtils] = PatternUtils): self.patterns: dict[str, list] = {} + self._pattern_utils: Final = pattern_utils def add_pattern(self, pattern: str, llm_deployment: dict): """ @@ -69,9 +70,10 @@ class PatternMatchRouter: """ # Convert the pattern to a regex regex: Final = self._pattern_to_regex(pattern) - if regex not in self.patterns: - self.patterns[regex] = [] - self.patterns[regex].append(llm_deployment) + if regex in self.patterns: + self.patterns[regex].append(llm_deployment) + return + self.patterns = dict(self._pattern_utils.sorted_patterns({**self.patterns, regex: [llm_deployment]})) def remove_deployment(self, model_id: str) -> None: """ @@ -138,11 +140,12 @@ class PatternMatchRouter: if request is None: return None - sorted_patterns: Final = PatternUtils.sorted_patterns(self.patterns) regex_filtered_model_names: Final = ( - [self._pattern_to_regex(m) for m in filtered_model_names] if filtered_model_names is not None else [] + tuple(self._pattern_to_regex(m) for m in filtered_model_names) + if filtered_model_names is not None + else () ) - for pattern, llm_deployments in sorted_patterns: + for pattern, llm_deployments in self.patterns.items(): if filtered_model_names is not None and pattern not in regex_filtered_model_names: continue pattern_match = re.match(pattern, request) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 31fb4fb55c5..a58d8125010 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6077,6 +6077,34 @@ def test_failure_handler_helper_fn_builds_payload_once_per_exception(): assert obj.model_call_details["standard_logging_object"] is not first_payload +@pytest.mark.asyncio +async def test_sync_failure_handler_reuses_payload_after_callable_async_callback(): + """Regression for LIT-6886: the proxy runs async_failure_handler, then the threaded + failure_handler, for every rejected request. A plain-function async callback (the + Router registers one) is dispatched through CustomLogger.async_log_event, which + restamps log_event_type on the shared model_call_details; the sync handler then + rebuilt the standardized payload, doubling the redaction and payload cost of a 403.""" + router_style_callback = AsyncMock() + obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6886-1", + function_id="f", + dynamic_async_failure_callbacks=[router_style_callback], + ) + exc = _raise_and_catch(_ClientError(status_code=403, message="key not allowed to access model")) + await obj.async_failure_handler(exception=exc, traceback_exception="") + first_payload = obj.model_call_details["standard_logging_object"] + assert first_payload is not None + assert router_style_callback.await_count == 1 + + obj.failure_handler(exc, "") + assert obj.model_call_details["standard_logging_object"] is first_payload + + @pytest.mark.asyncio async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_obj): """The savings gate reads litellm_gateway_injected_cache from the request's diff --git a/tests/test_litellm/router_utils/test_pattern_match_deployments.py b/tests/test_litellm/router_utils/test_pattern_match_deployments.py index 795d448ef5f..f9d9345cd26 100644 --- a/tests/test_litellm/router_utils/test_pattern_match_deployments.py +++ b/tests/test_litellm/router_utils/test_pattern_match_deployments.py @@ -2,8 +2,10 @@ from __future__ import annotations +from unittest.mock import Mock + from litellm.router_utils import pattern_match_deployments -from litellm.router_utils.pattern_match_deployments import PatternMatchRouter +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils def _wildcard_deployment(model_name: str) -> dict: @@ -76,3 +78,31 @@ def test_get_pattern_still_resolves_unqualified_names(monkeypatch): router = PatternMatchRouter() router.add_pattern("openai/*", _wildcard_deployment("openai/*")) assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"] + + +class _CountingPatternUtils(PatternUtils): + sorted_patterns = staticmethod(Mock(wraps=PatternUtils.sorted_patterns)) + + +def test_route_never_sorts_and_the_most_specific_pattern_still_wins_after_registry_changes(): + """Regression for LIT-6886: the auth layer walks the wildcard registry for every request, so an + unmatched model name (an invalid-model 403) re-sorted every pattern by specificity per request and + a burst of rejections saturated the worker CPU. Lookups must not sort; adding a pattern or removing + a deployment must still leave the most specific pattern winning.""" + router = PatternMatchRouter(pattern_utils=_CountingPatternUtils) + router.add_pattern("openai/*", _wildcard_deployment("openai/*")) + router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*")) + router.add_pattern("openai/gpt-*", {"model_name": "openai/gpt-*", "litellm_params": {"model": "azure/gpt-*"}}) + sorts_after_setup = _CountingPatternUtils.sorted_patterns.call_count + + for _ in range(3): + assert router.route("does-not-exist") is None + assert _matched_models(router.route("openai/gpt-4o")) == ["azure/gpt-4o"] + assert _matched_models(router.route("openai/o3")) == ["openai/o3"] + assert _CountingPatternUtils.sorted_patterns.call_count == sorts_after_setup + + router.add_pattern("openai/*", {**_wildcard_deployment("openai/*"), "model_info": {"id": "id-1"}}) + assert len(_matched_models(router.route("openai/o3"))) == 2 + router.remove_deployment("id-1") + assert _matched_models(router.route("openai/gpt-4o")) == ["azure/gpt-4o"] + assert _matched_models(router.route("openai/o3")) == ["openai/o3"] From 45cc2ed08225300536894028624284e6f9eb8baf Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 11:53:22 -0700 Subject: [PATCH 52/61] test(e2e): require a 200 inside the regenerate grace window and drop the helper docstrings --- tests/e2e/access_control/test_access_control_e2e.py | 6 +++--- tests/e2e/management/test_key_management_e2e.py | 4 ---- tests/e2e/management/test_management_e2e.py | 5 +++-- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/e2e/access_control/test_access_control_e2e.py b/tests/e2e/access_control/test_access_control_e2e.py index c30dadc49ae..9d01f2915e7 100644 --- a/tests/e2e/access_control/test_access_control_e2e.py +++ b/tests/e2e/access_control/test_access_control_e2e.py @@ -76,8 +76,6 @@ class TestAccessControl: def test_llm_api_routes_group_grants_every_llm_endpoint( self, client: AccessControlClient, resources: ResourceManager ) -> None: - """allowed_routes=["llm_api_routes"] names a route group, not a path: one - entry must open every LLM endpoint while the management routes stay shut.""" key = client.llm_only_key() resources.defer(lambda: client.delete_key(key)) @@ -89,7 +87,9 @@ class TestAccessControl: f"200 must carry a real completion, not an error envelope: {chat.body[:300]}" ) - embedding = unwrap(client.proxy.embed(key, EmbedBody(model=EMBEDDING_MODEL, input=f"route group {unique_marker()}"))) + embedding = unwrap( + client.proxy.embed(key, EmbedBody(model=EMBEDDING_MODEL, input=f"route group {unique_marker()}")) + ) assert embedding.model, f"llm_api_routes key reached /embeddings but got no model back: {embedding}" denied = client.create_model_status(key, f"e2e-route-group-{unique_marker()}") diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index d4347b8c0e7..8b7d5f0eb6f 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -90,8 +90,6 @@ def _is_budget_block(outcome: StreamingResponse) -> bool: def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None: - """Drive paid calls until the key's max_budget refuses one. The first call spends, - the reservation counter trips the cap, and the next call is the 429.""" for _ in range(40): outcome = client.chat_status(key, SPEND_MODEL, f"spend {unique_marker()}") if _is_budget_block(outcome): @@ -105,8 +103,6 @@ def _spend_until_budget_blocks(client: ManagementClient, key: str) -> None: def _settled_spend(client: ManagementClient, key: str) -> float | None: - """The key's recorded spend once it is positive and unchanged across two reads a - poll interval apart, so no batched spend write is still in flight when we reset.""" first = client.proxy.key_info(key).spend or 0.0 time.sleep(client.proxy.poll_interval) second = client.proxy.key_info(key).spend or 0.0 diff --git a/tests/e2e/management/test_management_e2e.py b/tests/e2e/management/test_management_e2e.py index eace8f2e3b4..476165b715d 100644 --- a/tests/e2e/management/test_management_e2e.py +++ b/tests/e2e/management/test_management_e2e.py @@ -12,6 +12,7 @@ from __future__ import annotations import math import time from collections.abc import Callable +from typing import Final import pytest @@ -377,12 +378,12 @@ class TestKeyRegeneration: new_key = client.regenerate_key(old_key, grace_period=REGENERATE_GRACE_PERIOD) resources.defer(lambda: client.proxy.delete_key(new_key)) - revoke_at = time.monotonic() + REGENERATE_GRACE_SECONDS + revoke_at: Final = time.monotonic() + REGENERATE_GRACE_SECONDS assert new_key != old_key, "regenerate returned the same key string, so no rotation happened" def old_accepted() -> bool | None: outcome = client.chat_status(old_key, "gpt-5.5", f"say hi {unique_marker()}") - return True if outcome.status_code != 401 else None + return True if outcome.ok else None _ = _poll(client, old_accepted, "old key was rejected 401 inside its grace period at the deadline") assert time.monotonic() < revoke_at, ( From ee5d66d030824c34e2bd15a1aecd62581b57dc61 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 11:56:56 -0700 Subject: [PATCH 53/61] fix(ui): keep Back to Guardrails reachable when a ?guardrail= link is stale With the selection in the URL, a mistyped or deleted guardrail id lands on the info view's not-found branch, which rendered only the message and left no way back to the table short of editing the address bar. The not-found branch now shares the Back to Guardrails button with the loaded view --- .../_components/guardrail_info.test.tsx | 24 +++++++++++++++++++ .../guardrails/_components/guardrail_info.tsx | 21 ++++++++-------- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index fcffc2122e7..836921104ff 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -338,3 +338,27 @@ describe("Guardrail Info", () => { expect(screen.getByText("Guardrail Settings")).toBeInTheDocument(); }); }); + +describe("Guardrail Info when the guardrail cannot be loaded", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should keep Back to Guardrails reachable so a stale ?guardrail= link is not a dead end", async () => { + vi.mocked(networking.getGuardrailInfo).mockRejectedValue(new Error("Guardrail stale-id not found")); + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: [], + }); + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + const onClose = vi.fn(); + + render(); + + expect(await screen.findByText("Guardrail not found")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: /back to guardrails/i })); + expect(onClose).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index aaa656d15bb..c9162d99934 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -481,16 +481,18 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, return
Loading...
; } + const backButton = ( + + ); + if (!guardrailData) { - return
Guardrail not found
; + return
{backButton}Guardrail not found
; } - // Format date helper function - const formatDate = (dateString?: string) => { - if (!dateString) return "-"; - const date = new Date(dateString); - return date.toLocaleString(); - }; + const formatDate = (dateString?: string) => (dateString ? new Date(dateString).toLocaleString() : "-"); // Format the provider display name and logo const { logo, displayName } = getGuardrailLogoAndName(guardrailData.litellm_params?.guardrail || ""); @@ -510,10 +512,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, return (
- + {backButton}

{guardrailData.guardrail_name || "Unnamed Guardrail"}

{guardrailData.guardrail_id}

From da5af0cb27aa668527a0e5746e307ab3c1188a24 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 12:02:48 -0700 Subject: [PATCH 54/61] test: repair two CI tests broken by intentional changes test_no_linear_scans_in_router: #39674 renamed heuristic_v2_router_limit_violation to auto_router_capability_violation, so the allowlist entry stopped matching and the same admin-only scan tripped the static check. Rename the entry to follow it. tableScrolling.spec.ts: 9ba6cab889 (LIT-4738) gave the Tags and Model Hub tables client-side pagination at 25 rows, so the 40 seeded rows no longer render on one page. Select 50 rows per page before counting, as the Logs case already does. --- tests/e2e/ui/tests/tables/tableScrolling.spec.ts | 2 ++ tests/router_unit_tests/test_router_index_management.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts index 5c7438cfd11..0537ba193be 100644 --- a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts +++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts @@ -184,6 +184,7 @@ test.describe("Admin tables scroll inside the page", () => { ); try { await navigateToPage(page, Page.TagManagement); + await setRowsPerPage(page, "50"); await expectRowsAtLeast(page, SEED_ROWS); expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); } finally { @@ -207,6 +208,7 @@ test.describe("Admin tables scroll inside the page", () => { ); try { await navigateToPage(page, Page.ModelHubTable); + await setRowsPerPage(page, "50"); await expectRowsAtLeast(page, SEED_ROWS); expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); } finally { diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py index 35d295d581a..291badd91b4 100644 --- a/tests/router_unit_tests/test_router_index_management.py +++ b/tests/router_unit_tests/test_router_index_management.py @@ -241,7 +241,7 @@ class TestRouterIndexManagement: "_get_deployment_by_litellm_model": "lookup by litellm_params.model, which is not indexed", "_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match', "config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)", - "heuristic_v2_router_limit_violation": "counts heuristic_v2 routers across the whole list; admin path only (auto-router init/upsert)", + "auto_router_capability_violation": "counts gated auto-routers across the whole list; admin path only (auto-router init/upsert)", } # Get path to router.py From 0ad361a7283498e5f8b0154486e1b9a2a97270cd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:02:55 +0000 Subject: [PATCH 55/61] fix(router): coordinate async and sync failure handlers at remaining router call sites (#39887) * fix(router): coordinate async and sync failure handlers at remaining router call sites Five router failure paths still scheduled logging_obj.async_failure_handler as a task while starting logging_obj.failure_handler on a raw thread, so both handlers mutated the same logging object concurrently. Route them through dispatch_failure_handlers like the streaming paths already do. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): wait on the real logging executor and justify the callbacks global patch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): submit sync failure handler even when the dispatch task is cancelled Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): justify the executor submit patch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 15 +- litellm/router.py | 53 +++---- .../test_litellm_logging.py | 56 ++++++++ tests/test_litellm/test_router.py | 132 ++++++++++++++++++ 4 files changed, 216 insertions(+), 40 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 22e4dbf3a44..83e0b4d84f1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1918,7 +1918,9 @@ class Logging(LiteLLMLoggingBaseClass): two paths cannot mutate it at the same time. ``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from ``completion()``); legacy string callbacks still run via - ``executor.submit(failure_handler)`` when configured. + ``executor.submit(failure_handler)`` when configured, and still get submitted + when the awaiting task is cancelled (e.g. the event loop shuts down right after + the request failed). """ litellm_params: Final = self.model_call_details.get("litellm_params", {}) or {} sync_sdk: Final = self._is_sync_litellm_request(litellm_params) @@ -1927,12 +1929,11 @@ class Logging(LiteLLMLoggingBaseClass): self.failure_handler(exception, traceback_exception) return - await self.async_failure_handler(exception, traceback_exception) - - if not self._should_run_sync_failure_callbacks_for_async_calls(): - return - - executor.submit(self.failure_handler, exception, traceback_exception) + try: + await self.async_failure_handler(exception, traceback_exception) + finally: + if self._should_run_sync_failure_callbacks_for_async_calls(): + executor.submit(self.failure_handler, exception, traceback_exception) def should_run_logging( self, diff --git a/litellm/router.py b/litellm/router.py index 490836f5f0a..76da3a857df 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8380,17 +8380,12 @@ class Router: ## LOG FAILURE EVENT if logging_obj is not None: asyncio.create_task( - logging_obj.async_failure_handler( + logging_obj.dispatch_failure_handlers( exception=e, traceback_exception=traceback.format_exc(), - end_time=time.time(), + prefer_async_handlers=True, ) ) - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback.format_exc()), - ).start() # log response _set_cooldown_deployments( litellm_router_instance=self, exception_status=e.status_code, @@ -8403,17 +8398,12 @@ class Router: ## LOG FAILURE EVENT if logging_obj is not None: asyncio.create_task( - logging_obj.async_failure_handler( + logging_obj.dispatch_failure_handlers( exception=e, traceback_exception=traceback.format_exc(), - end_time=time.time(), + prefer_async_handlers=True, ) ) - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback.format_exc()), - ).start() # log response raise e async def async_callback_filter_deployments( @@ -8451,17 +8441,12 @@ class Router: ## LOG FAILURE EVENT if logging_obj is not None: asyncio.create_task( - logging_obj.async_failure_handler( + logging_obj.dispatch_failure_handlers( exception=e, traceback_exception=traceback.format_exc(), - end_time=time.time(), + prefer_async_handlers=True, ) ) - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback.format_exc()), - ).start() # log response raise e return returned_healthy_deployments @@ -12643,13 +12628,13 @@ class Router: logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) if logging_obj is not None: - ## LOGGING - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming - asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception=e, + traceback_exception=traceback_exception, + prefer_async_handlers=True, + ) + ) raise e async def async_get_available_deployment_for_pass_through( @@ -12777,11 +12762,13 @@ class Router: if request_kwargs is not None: logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) if logging_obj is not None: - threading.Thread( - target=logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() - asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception=e, + traceback_exception=traceback_exception, + prefer_async_handlers=True, + ) + ) raise e async def _run_routing_plugins( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index a58d8125010..1991170707d 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1546,6 +1546,62 @@ async def test_dispatch_failure_handlers_async_completes_before_sync_submit( assert events == ["async_start", "async_end", "sync_submit"] +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_submits_sync_handler_when_task_is_cancelled( + logging_obj, +): + """Cancelling the dispatch task mid-await still submits the sync failure_handler. + + Router failure paths fire the dispatcher with ``asyncio.create_task`` and raise + right away. When the event loop is torn down before the task finishes (a short + ``asyncio.run`` in the SDK), the cancelled task must still hand the sync callbacks + to the executor, as the old raw-thread path did, and only once the async handler + has stopped. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + events: list[str] = [] + async_started = asyncio.Event() + + async def _async_failure(exc, tb, **kwargs): + events.append("async_start") + async_started.set() + await asyncio.sleep(10) + events.append("async_end") + + def _submit(*args, **kwargs): + events.append("sync_submit") + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", new_callable=MagicMock), + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=True, + ), + patch( # test-quality-ok: the executor submit is the observable + "litellm.litellm_core_utils.litellm_logging.executor.submit", + side_effect=_submit, + ), + ): + task = asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + ) + await async_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert events == ["async_start", "sync_submit"] + + @pytest.mark.asyncio async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks( logging_obj, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ffd6c5f97ce..8368aa11316 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5,6 +5,7 @@ import json import logging import os import threading +from datetime import datetime from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -20,6 +21,7 @@ import litellm from litellm import Router from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, @@ -12940,6 +12942,136 @@ async def test_router_retry_policy_controls_upstream_attempt_count( assert upstream.call_count == expected_upstream_calls +def _make_failure_logging_obj(): + return LiteLLMLogging( + model="gpt-5.6", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="lit-6960", + function_id="f", + ) + + +async def _assert_router_failure_logging_is_coordinated(logging_obj, trigger, expected_exception): + """The sync failure_handler must not start until async_failure_handler has finished on the shared logging_obj.""" + events: list[str] = [] + sync_done = threading.Event() + + async def _async_failure(*args, **kwargs): + events.append("async_start") + await asyncio.sleep(0.05) + events.append("async_end") + + def _sync_failure(*args, **kwargs): + events.append("sync_start") + sync_done.set() + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", side_effect=_sync_failure), + patch.object(logging_obj, "_should_run_sync_failure_callbacks_for_async_calls", return_value=True), + ): + with pytest.raises(expected_exception): + await trigger() + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + await asyncio.gather(*pending) + assert await asyncio.to_thread(sync_done.wait, 5), "failure_handler never ran" + + assert events == ["async_start", "async_end", "sync_start"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "hook_error", + [ + litellm.RateLimitError(message="rpm exceeded", llm_provider="openai", model="gpt-5.6"), + RuntimeError("pre call check blew up"), + ], +) +async def test_async_routing_strategy_pre_call_checks_failure_logging_is_coordinated(hook_error): + class _RaisingPreCallCheck(CustomLogger): + async def async_pre_call_check(self, deployment, parent_otel_span): + raise hook_error + + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + deployment = router.model_list[0] + logging_obj = _make_failure_logging_obj() + + with patch.object(litellm, "callbacks", [_RaisingPreCallCheck()]): # test-quality-ok: router reads this global + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_routing_strategy_pre_call_checks( + deployment=deployment, parent_otel_span=None, logging_obj=logging_obj + ), + type(hook_error), + ) + + +@pytest.mark.asyncio +async def test_async_callback_filter_deployments_failure_logging_is_coordinated(): + class _RaisingFilter(CustomLogger): + async def async_filter_deployments(self, *args, **kwargs): + raise RuntimeError("filter blew up") + + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + logging_obj = _make_failure_logging_obj() + + with patch.object(litellm, "callbacks", [_RaisingFilter()]): # test-quality-ok: router reads this global + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_callback_filter_deployments( + model="gpt-5.6", + healthy_deployments=router.model_list, + messages=None, + parent_otel_span=None, + request_kwargs={}, + logging_obj=logging_obj, + ), + RuntimeError, + ) + + +@pytest.mark.asyncio +async def test_async_get_available_deployment_failure_logging_is_coordinated(): + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + logging_obj = _make_failure_logging_obj() + + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_get_available_deployment( + model="model-that-is-not-configured", + request_kwargs={"litellm_logging_obj": logging_obj}, + messages=[{"role": "user", "content": "hi"}], + ), + litellm.BadRequestError, + ) + + +@pytest.mark.asyncio +async def test_async_get_available_deployment_for_pass_through_failure_logging_is_coordinated(): + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}] + ) + logging_obj = _make_failure_logging_obj() + + await _assert_router_failure_logging_is_coordinated( + logging_obj, + lambda: router.async_get_available_deployment_for_pass_through( + model="gpt-5.6", + request_kwargs={"litellm_logging_obj": logging_obj}, + ), + litellm.BadRequestError, + ) + + class _InFlightTracker: def __init__(self) -> None: self.current = 0 From 5298deb491ca1485f79a93f337e665421ec42b2e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 12:04:56 -0700 Subject: [PATCH 56/61] test(e2e/ui): select 50 rows per page before asserting the Tags and Model Hub tables overflow #39680 made every admin table honor the selected page size, so the Tags and Model Hub tables now paginate at 25 by default and the two scroll specs, which seed 40 rows and expect them all on one page, fail on every litellm-e2e-ui run since (builds 206 to 208). Selecting 50 rows first, the way the Request Logs spec already does, keeps the overflow assertion meaningful --- tests/e2e/ui/tests/tables/tableScrolling.spec.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts index 5c7438cfd11..0537ba193be 100644 --- a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts +++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts @@ -184,6 +184,7 @@ test.describe("Admin tables scroll inside the page", () => { ); try { await navigateToPage(page, Page.TagManagement); + await setRowsPerPage(page, "50"); await expectRowsAtLeast(page, SEED_ROWS); expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); } finally { @@ -207,6 +208,7 @@ test.describe("Admin tables scroll inside the page", () => { ); try { await navigateToPage(page, Page.ModelHubTable); + await setRowsPerPage(page, "50"); await expectRowsAtLeast(page, SEED_ROWS); expect(await rowsPaintingPastAnAncestor(page)).toEqual([]); } finally { From 73e1cfb378e9d45c0a92266d6a763e61440cc862 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 5 Sep 2026 12:09:53 -0700 Subject: [PATCH 57/61] fix(cloudzero): infer daily batch schema from every row (#39871) * fix(cloudzero): infer daily batch schema from every row pl.DataFrame defaults to inferring column types from the first 100 rows, so a day whose batch starts with more than 100 rows missing team_alias, api_key_alias or user_email typed that column as Null and then raised a ComputeError on the first row that had a value, failing the whole export with a 500 and sending nothing. Pass infer_schema_length=None when rebuilding each day's DataFrame, the same guard the usage query already uses. * test(cloudzero): cover late tag schema inference Exercise the CloudZero resource tag field after a long run of missing values so a finite inference window fails the regression test. --- .../integrations/cloudzero/cz_stream_api.py | 6 ++++- .../cloudzero/test_cz_stream_api.py | 24 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 1e2fa318786..2213c5fe275 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -97,7 +97,11 @@ class CloudZeroStreamer: continue # Convert lists back to DataFrames - return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records} + return { + date_key: pl.DataFrame(records, infer_schema_length=None) + for date_key, records in daily_batches.items() + if records + } def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime: """Parse timestamp string and convert to UTC.""" diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py index 1a95e45b2d5..d4e49a1252f 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py +++ b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py @@ -69,6 +69,30 @@ class TestCloudZeroStreamer: assert "2025-01-19" in result assert len(result["2025-01-19"]) == 1 + def test_group_by_date_infers_schema_from_every_row(self): + """Test daily batches retain optional string columns that are null for thousands of leading rows.""" + streamer = CloudZeroStreamer("test-key", "test-connection") + leading_nulls = 10_000 + rows = [ + {"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": None} + for _ in range(leading_nulls) + ] + rows.append( + {"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": "team-alias"} + ) + data = pl.DataFrame( + rows, + schema={"time/usage_start": pl.String, "resource/tag:team_alias": pl.String}, + ) + + result = streamer._group_by_date(data) + + batch = result["2025-01-19"] + assert len(batch) == leading_nulls + 1 + assert batch.schema["resource/tag:team_alias"] == pl.String + assert batch["resource/tag:team_alias"].null_count() == leading_nulls + assert batch.tail(1).item(0, "resource/tag:team_alias") == "team-alias" + def test_parse_and_convert_timestamp_utc(self): """Test _parse_and_convert_timestamp method with UTC timestamp.""" streamer = CloudZeroStreamer("test-key", "test-connection") From 877197918bfe7540e714c6ef2acfb24694df5049 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 5 Sep 2026 12:10:05 -0700 Subject: [PATCH 58/61] fix(cloudzero): preserve late resource tags (#39873) * fix(cloudzero): infer daily batch schema from every row pl.DataFrame defaults to inferring column types from the first 100 rows, so a day whose batch starts with more than 100 rows missing team_alias, api_key_alias or user_email typed that column as Null and then raised a ComputeError on the first row that had a value, failing the whole export with a 500 and sending nothing. Pass infer_schema_length=None when rebuilding each day's DataFrame, the same guard the usage query already uses. * test(cloudzero): cover late tag schema inference Exercise the CloudZero resource tag field after a long run of missing values so a finite inference window fails the regression test. * fix(cloudzero): preserve late resource tags * style(cloudzero): remove redundant test comment --- litellm/integrations/cloudzero/transform.py | 2 +- .../integrations/cloudzero/test_transform.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index ffc8fe1c1f5..12a0ee55fad 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -95,7 +95,7 @@ class CBFTransformer: if len(cbf_data) > 0: console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]") - return pl.DataFrame(cbf_data) + return pl.DataFrame(cbf_data, infer_schema_length=None) def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord: """Create a single CBF record from LiteLLM daily spend row.""" diff --git a/tests/test_litellm/integrations/cloudzero/test_transform.py b/tests/test_litellm/integrations/cloudzero/test_transform.py index 3ec2fe6779e..cf8d70702f9 100644 --- a/tests/test_litellm/integrations/cloudzero/test_transform.py +++ b/tests/test_litellm/integrations/cloudzero/test_transform.py @@ -86,6 +86,33 @@ class TestCBFTransformer: assert result.is_empty() + def test_transform_keeps_tags_first_seen_after_row_100(self): + transformer = CBFTransformer() + teamless_rows = 101 + team_rows = 2 + total_rows = teamless_rows + team_rows + data = pl.DataFrame( + { + "date": ["2025-01-19"] * total_rows, + "successful_requests": [1] * total_rows, + "spend": [0.5] * total_rows, + "prompt_tokens": [10] * total_rows, + "completion_tokens": [5] * total_rows, + "model": ["gpt-4"] * total_rows, + "custom_llm_provider": ["openai"] * total_rows, + "api_key": ["sk-late-team"] * total_rows, + "team_id": pl.Series([None] * teamless_rows + ["team-late"] * team_rows, dtype=pl.String), + "team_alias": pl.Series([None] * teamless_rows + ["Late Team"] * team_rows, dtype=pl.String), + } + ) + + result = transformer.transform(data) + + assert len(result) == total_rows + assert "resource/tag:team_alias" in result.columns + assert result["resource/tag:team_alias"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows + assert result["resource/tag:entity_id"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows + def test_create_cbf_record(self): """Test _create_cbf_record method with valid row data.""" transformer = CBFTransformer() From b290dd410e6bb9c59fc3ac7219a3bb197cf027d8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:22:52 -0700 Subject: [PATCH 59/61] feat(terraform/gcp): dependencies-only mode and bring-your-own-network for GKE (#39695) * feat(terraform/gcp): add dependencies-only mode Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform/gcp): review fixes for dependencies-only mode Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-terraform-modules.yml | 31 ++++ terraform/litellm/gcp/README.md | 77 +++++++- terraform/litellm/gcp/bootstrap.tf | 6 +- terraform/litellm/gcp/cloudrun.tf | 109 +++++------ terraform/litellm/gcp/cloudsql.tf | 9 +- .../litellm/gcp/examples/default/main.tf | 5 + .../litellm/gcp/examples/default/outputs.tf | 30 +++ .../examples/default/terraform.tfvars.example | 8 + .../litellm/gcp/examples/default/variables.tf | 24 +++ terraform/litellm/gcp/iam.tf | 11 ++ terraform/litellm/gcp/load_balancer.tf | 58 ++++-- terraform/litellm/gcp/locals.tf | 5 +- terraform/litellm/gcp/network.tf | 20 +- terraform/litellm/gcp/outputs.tf | 58 ++++-- terraform/litellm/gcp/redis.tf | 4 +- .../litellm/gcp/tests/deps_only.tftest.hcl | 175 ++++++++++++++++++ terraform/litellm/gcp/variables.tf | 30 ++- 17 files changed, 540 insertions(+), 120 deletions(-) create mode 100644 terraform/litellm/gcp/tests/deps_only.tftest.hcl diff --git a/.github/workflows/test-terraform-modules.yml b/.github/workflows/test-terraform-modules.yml index 0e3e5330453..52006d9b578 100644 --- a/.github/workflows/test-terraform-modules.yml +++ b/.github/workflows/test-terraform-modules.yml @@ -4,6 +4,7 @@ on: push: paths: - "terraform/litellm/aws/**" + - "terraform/litellm/gcp/**" - ".github/workflows/test-terraform-modules.yml" pull_request: branches: @@ -13,6 +14,7 @@ on: - "litellm_**" paths: - "terraform/litellm/aws/**" + - "terraform/litellm/gcp/**" - ".github/workflows/test-terraform-modules.yml" permissions: @@ -52,3 +54,32 @@ jobs: # Plan-only, mock_provider-backed: no AWS credentials, no API calls. - name: test run: terraform test + + gcp-module: + name: fmt, validate, test (gcp) + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: terraform/litellm/gcp + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2 + with: + terraform_version: 1.13.3 + terraform_wrapper: false + + - name: fmt + run: terraform fmt -recursive -check -diff + + - name: init + run: terraform init -backend=false -input=false + + - name: validate + run: terraform validate + + - name: test + run: terraform test diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 88e9979148f..c93e5f6b303 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -392,6 +392,63 @@ with its own provider config (one `examples/default`-style root per project), or fork the module to add `configuration_aliases` and pass per-instance `providers = { ... }`. +## Dependencies only (run LiteLLM on GKE) + +Set `create_runtime = false` to provision Cloud SQL, Memorystore, GCS, +Secret Manager, and the runtime service account without Cloud Run or the +load balancer. For a Shared VPC, set the full host-project network ID and +skip PSA creation after the host project has configured it: + +```hcl +create_runtime = false +network_id = "projects//global/networks/" +create_psa_connection = false +``` + +The host project must already have Private Services Access configured on +that network and the Service Networking API enabled; the module cannot set +PSA up from a service project. GKE nodes must sit on the same Shared VPC so +the Cloud SQL and Memorystore private IPs are routable from the pods. Run +the root with its provider pointed at the project that should own the +dependencies. `create_runtime = true` with `network_id` set is also allowed, +but the Serverless VPC Access connector has to live in the same project as +the network, so that combination only works when the VPC is in the +deployment project + +Map the outputs into the Helm values as follows: + +```yaml +database: + writer: + host: + dbname: + passwordSecret: + name: + reader: + host: + dbname: + passwordSecret: + name: +redis: + host: + port: +masterKey: + secretName: +``` + +Create the database Secret with keys `username` (the `db_username` output) +and `password` (read it with `gcloud secrets versions access latest +--secret=`), and the master key Secret from +`master_key_secret_id` the same way. Memorystore only accepts TLS by +default, so store the `redis_server_ca_pem` output in a third Secret, +mount it into the gateway and backend pods via `volumes` / `volumeMounts`, +and add `REDIS_SSL=true` and `REDIS_SSL_CA_CERTS=` to each +component's `extraEnv`. Setting `redis_transit_encryption = false` removes +the CA plumbing at the cost of plaintext Redis traffic inside the VPC + +The chart's pre-install/pre-upgrade migration hook runs the Prisma +migration, so nothing replaces the Cloud Run migrations Job in this mode + ## Storage and database retention Two opt-in tripwires guard against accidental data loss on @@ -409,14 +466,15 @@ Flip `cloudsql_deletion_protection` to `false` or `gcs_force_destroy` to ## Redis encryption -Memorystore runs with `transit_encryption_mode = "SERVER_AUTHENTICATION"`, -so the proxy connects via `rediss://`. The instance's self-signed CA cert -(`server_ca_certs[0].cert`) is shipped to gateway + backend as -`REDIS_CA_PEM_B64`; their entrypoint shell decodes it to `/tmp/redis-ca.pem` -before uvicorn starts and points `REDIS_SSL_CA_CERTS` at that path. No -extra config needed — but if you ever swap Memorystore for an external -Redis, override `REDIS_HOST`/`REDIS_PORT` and either drop these env vars -or point them at your own CA. +By default, Memorystore runs with +`transit_encryption_mode = "SERVER_AUTHENTICATION"`, so Cloud Run connects +via `rediss://`. The instance's self-signed CA cert +(`server_ca_certs[0].cert`) is shipped to gateway and backend as +`REDIS_CA_PEM_B64`; their entrypoint shell decodes it to +`/tmp/redis-ca.pem` before uvicorn starts and points `REDIS_SSL_CA_CERTS` at +that path. Set `redis_transit_encryption = false` to use plaintext Redis. +For GKE, use `redis_server_ca_pem` as described in the dependencies-only +section, or accept the security tradeoff of disabling transit encryption ## Files @@ -434,4 +492,5 @@ or point them at your own CA. | `iam.tf` | Runtime SA + Cloud SQL client + Secret Manager accessor | | `cloudrun.tf` | 3 Cloud Run services + Cloud Run Job for migrations | | `load_balancer.tf`| External HTTPS LB, serverless NEGs, URL map for path routing | -| `outputs.tf` | LB IP, service URLs, secret IDs, migration `execute` command | +| `outputs.tf` | LB IP, service URLs, dependency endpoints, secret IDs, migration command | +| `tests/` | Plan-only mock-provider coverage for deployment modes and Redis encryption | diff --git a/terraform/litellm/gcp/bootstrap.tf b/terraform/litellm/gcp/bootstrap.tf index b929c4d76f3..dead5c41f6b 100644 --- a/terraform/litellm/gcp/bootstrap.tf +++ b/terraform/litellm/gcp/bootstrap.tf @@ -15,15 +15,17 @@ # enough to invoke Cloud Run admin APIs (`gcloud auth login`). resource "terraform_data" "migration" { + count = var.create_runtime ? 1 : 0 + triggers_replace = { - job_id = google_cloud_run_v2_job.migrations.id + job_id = google_cloud_run_v2_job.migrations[0].id job_image = local.migrations_image } provisioner "local-exec" { interpreter = ["bash", "-c"] environment = { - JOB = google_cloud_run_v2_job.migrations.name + JOB = google_cloud_run_v2_job.migrations[0].name REGION = var.region PROJECT = var.project_id } diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 5a5c361b832..84ae8b9247f 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -6,25 +6,28 @@ locals { # Memorystore exposes a self-signed CA cert per instance; we ship it as # a base64 env var and decode it to a file at container startup so the # rediss:// connection can validate. Public cert, not sensitive. - redis_ca_pem_b64 = base64encode(google_redis_instance.this.server_ca_certs[0].cert) + redis_ca_pem_b64 = var.redis_transit_encryption ? base64encode(google_redis_instance.this.server_ca_certs[0].cert) : "" - shared_env_kv = [ - { name = "DATABASE_HOST", value = google_sql_database_instance.writer.private_ip_address }, - { name = "DATABASE_PORT", value = "5432" }, - { name = "DATABASE_USER", value = var.db_username }, - { name = "DATABASE_NAME", value = var.db_name }, - { name = "DATABASE_HOST_READ_REPLICA", value = google_sql_database_instance.reader.private_ip_address }, - { name = "DATABASE_PORT_READ_REPLICA", value = "5432" }, - { name = "REDIS_HOST", value = google_redis_instance.this.host }, - { name = "REDIS_PORT", value = tostring(google_redis_instance.this.port) }, - # _redis.get_redis_url_from_environment honors REDIS_SSL to flip the - # scheme to rediss://; REDIS_SSL_CA_CERTS is mapped via - # _get_redis_env_kwarg_mapping → ssl_ca_certs on the redis-py client. - { name = "REDIS_SSL", value = "true" }, - { name = "REDIS_SSL_CA_CERTS", value = "/tmp/redis-ca.pem" }, - { name = "REDIS_CA_PEM_B64", value = local.redis_ca_pem_b64 }, - { name = "GCS_BUCKET_NAME", value = google_storage_bucket.this.name }, - ] + shared_env_kv = concat( + [ + { name = "DATABASE_HOST", value = google_sql_database_instance.writer.private_ip_address }, + { name = "DATABASE_PORT", value = "5432" }, + { name = "DATABASE_USER", value = var.db_username }, + { name = "DATABASE_NAME", value = var.db_name }, + { name = "DATABASE_HOST_READ_REPLICA", value = google_sql_database_instance.reader.private_ip_address }, + { name = "DATABASE_PORT_READ_REPLICA", value = "5432" }, + { name = "REDIS_HOST", value = google_redis_instance.this.host }, + { name = "REDIS_PORT", value = tostring(google_redis_instance.this.port) }, + ], + var.redis_transit_encryption ? [ + { name = "REDIS_SSL", value = "true" }, + { name = "REDIS_SSL_CA_CERTS", value = "/tmp/redis-ca.pem" }, + { name = "REDIS_CA_PEM_B64", value = local.redis_ca_pem_b64 }, + ] : [], + [ + { name = "GCS_BUCKET_NAME", value = google_storage_bucket.this.name }, + ], + ) # OTel v2 is opt-in and gated on otel_endpoint, matching the AWS stack — # nothing OTel-related is added to the container env until an endpoint is @@ -126,9 +129,9 @@ locals { # Decode the Memorystore CA cert (passed as REDIS_CA_PEM_B64) to the # path REDIS_SSL_CA_CERTS points at, so the redis-py client can validate # the rediss:// handshake. - redis_ca_fragment = [ + redis_ca_fragment = var.redis_transit_encryption ? [ "python -c \"import os, base64, pathlib; pathlib.Path(os.environ['REDIS_SSL_CA_CERTS']).write_bytes(base64.b64decode(os.environ['REDIS_CA_PEM_B64']))\"" - ] + ] : [] database_url_fragment = [ "export DATABASE_URL=\"postgresql://$${DATABASE_USER}:$${DATABASE_PASSWORD}@$${DATABASE_HOST}:$${DATABASE_PORT}/$${DATABASE_NAME}\"", @@ -171,29 +174,7 @@ locals { # ---------- Gateway ---------- resource "google_cloud_run_v2_service" "gateway" { - # Metering needs a client certificate AND its key. Each secret is created only - # when its own PEM is supplied, so an endpoint set with a missing key would - # otherwise apply cleanly and leave the proxy logging "missing config" and - # never exporting. ca_cert_pem stays optional: empty means fall back to the - # system trust store. - # - # The guard lives here, on an unconditional resource, rather than on the cert - # secret: that secret is count-gated on the cert itself, so it has zero - # instances in exactly the case this must catch. Adding count or for_each to - # this resource would silently stop the guard from evaluating. - # - # endpoint cert key -> result - # "" any any -> metering off, no secrets created - # set set set -> metering on - # set any-missing -> plan fails here - lifecycle { - precondition { - condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" - ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." - } - } + count = var.create_runtime ? 1 : 0 name = "${local.name}-gateway" location = var.region @@ -206,7 +187,7 @@ resource "google_cloud_run_v2_service" "gateway" { max_instance_request_concurrency = var.gateway_max_instance_request_concurrency vpc_access { - connector = google_vpc_access_connector.this.id + connector = google_vpc_access_connector.this[0].id egress = "PRIVATE_RANGES_ONLY" } @@ -312,17 +293,7 @@ resource "google_cloud_run_v2_service" "gateway" { # ---------- Backend ---------- resource "google_cloud_run_v2_service" "backend" { - # Same guard as the gateway: the backend meters too (it serves the named-server - # MCP transport), and a targeted apply of just this resource must not slip a - # billing endpoint through without the credentials to use it. - lifecycle { - precondition { - condition = var.billing_metrics_endpoint == "" || ( - var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" - ) - error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set." - } - } + count = var.create_runtime ? 1 : 0 name = "${local.name}-backend" location = var.region @@ -335,7 +306,7 @@ resource "google_cloud_run_v2_service" "backend" { max_instance_request_concurrency = var.backend_max_instance_request_concurrency vpc_access { - connector = google_vpc_access_connector.this.id + connector = google_vpc_access_connector.this[0].id egress = "PRIVATE_RANGES_ONLY" } @@ -443,6 +414,8 @@ resource "google_cloud_run_v2_service" "backend" { # with zero IAM bindings, so a compromised UI container can't pivot to # Secret Manager / Cloud SQL via the metadata service. resource "google_cloud_run_v2_service" "ui" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-ui" location = var.region ingress = "INGRESS_TRAFFIC_INTERNAL_LOAD_BALANCER" @@ -450,7 +423,7 @@ resource "google_cloud_run_v2_service" "ui" { deletion_protection = false template { - service_account = google_service_account.ui_runtime.email + service_account = google_service_account.ui_runtime[0].email max_instance_request_concurrency = var.ui_max_instance_request_concurrency scaling { @@ -491,25 +464,31 @@ resource "google_cloud_run_v2_service" "ui" { # (LITELLM_MASTER_KEY); these IAM bindings just open up Cloud Run's invoker # gate so the LB request makes it to the container. resource "google_cloud_run_v2_service_iam_member" "gateway_allusers" { + count = var.create_runtime ? 1 : 0 + project = var.project_id - location = google_cloud_run_v2_service.gateway.location - name = google_cloud_run_v2_service.gateway.name + location = google_cloud_run_v2_service.gateway[0].location + name = google_cloud_run_v2_service.gateway[0].name role = "roles/run.invoker" member = "allUsers" } resource "google_cloud_run_v2_service_iam_member" "backend_allusers" { + count = var.create_runtime ? 1 : 0 + project = var.project_id - location = google_cloud_run_v2_service.backend.location - name = google_cloud_run_v2_service.backend.name + location = google_cloud_run_v2_service.backend[0].location + name = google_cloud_run_v2_service.backend[0].name role = "roles/run.invoker" member = "allUsers" } resource "google_cloud_run_v2_service_iam_member" "ui_allusers" { + count = var.create_runtime ? 1 : 0 + project = var.project_id - location = google_cloud_run_v2_service.ui.location - name = google_cloud_run_v2_service.ui.name + location = google_cloud_run_v2_service.ui[0].location + name = google_cloud_run_v2_service.ui[0].name role = "roles/run.invoker" member = "allUsers" } @@ -519,6 +498,8 @@ resource "google_cloud_run_v2_service_iam_member" "ui_allusers" { # assembles DATABASE_URL from the DATABASE_* env vars and runs `prisma # migrate deploy`. No proxy_config, no master key, no shell wrapper. resource "google_cloud_run_v2_job" "migrations" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-migrations" location = var.region labels = local.labels @@ -529,7 +510,7 @@ resource "google_cloud_run_v2_job" "migrations" { service_account = google_service_account.runtime.email vpc_access { - connector = google_vpc_access_connector.this.id + connector = google_vpc_access_connector.this[0].id egress = "PRIVATE_RANGES_ONLY" } diff --git a/terraform/litellm/gcp/cloudsql.tf b/terraform/litellm/gcp/cloudsql.tf index c9c2d03b2de..777434b4727 100644 --- a/terraform/litellm/gcp/cloudsql.tf +++ b/terraform/litellm/gcp/cloudsql.tf @@ -36,7 +36,7 @@ resource "google_sql_database_instance" "writer" { ip_configuration { ipv4_enabled = false - private_network = google_compute_network.this.id + private_network = local.network_id } insights_config { @@ -55,6 +55,11 @@ resource "google_sql_database_instance" "writer" { # (full data loss). Set the initial size only; let Cloud SQL own it # thereafter. ignore_changes = [settings[0].disk_size] + + precondition { + condition = var.create_psa_connection || var.network_id != "" + error_message = "create_psa_connection must be true unless network_id references an existing VPC with Private Services Access configured." + } } } @@ -76,7 +81,7 @@ resource "google_sql_database_instance" "reader" { ip_configuration { ipv4_enabled = false - private_network = google_compute_network.this.id + private_network = local.network_id } } diff --git a/terraform/litellm/gcp/examples/default/main.tf b/terraform/litellm/gcp/examples/default/main.tf index 8760d445f0c..f44b2a9a001 100644 --- a/terraform/litellm/gcp/examples/default/main.tf +++ b/terraform/litellm/gcp/examples/default/main.tf @@ -31,6 +31,11 @@ module "litellm" { tenant = var.tenant env = var.env + create_runtime = var.create_runtime + network_id = var.network_id + create_psa_connection = var.create_psa_connection + redis_transit_encryption = var.redis_transit_encryption + litellm_master_key = var.litellm_master_key litellm_license = var.litellm_license ui_password = var.ui_password diff --git a/terraform/litellm/gcp/examples/default/outputs.tf b/terraform/litellm/gcp/examples/default/outputs.tf index 3a9343c4850..48cdc1af66e 100644 --- a/terraform/litellm/gcp/examples/default/outputs.tf +++ b/terraform/litellm/gcp/examples/default/outputs.tf @@ -38,6 +38,31 @@ output "redis_endpoint" { value = module.litellm.redis_endpoint } +output "redis_host" { + description = "Memorystore Redis host." + value = module.litellm.redis_host +} + +output "redis_port" { + description = "Memorystore Redis port." + value = module.litellm.redis_port +} + +output "redis_server_ca_pem" { + description = "Memorystore server CA PEM." + value = module.litellm.redis_server_ca_pem +} + +output "db_username" { + description = "Cloud SQL application username." + value = module.litellm.db_username +} + +output "db_name" { + description = "Cloud SQL database name." + value = module.litellm.db_name +} + output "gcs_bucket" { description = "GCS bucket name." value = module.litellm.gcs_bucket @@ -53,6 +78,11 @@ output "db_password_secret_id" { value = module.litellm.db_password_secret_id } +output "runtime_service_account_email" { + description = "Runtime service account email." + value = module.litellm.runtime_service_account_email +} + output "migration_run_command" { description = "Break-glass command to re-run the one-off migration job." value = module.litellm.migration_run_command diff --git a/terraform/litellm/gcp/examples/default/terraform.tfvars.example b/terraform/litellm/gcp/examples/default/terraform.tfvars.example index 4416cf0ee5d..c35206503bb 100644 --- a/terraform/litellm/gcp/examples/default/terraform.tfvars.example +++ b/terraform/litellm/gcp/examples/default/terraform.tfvars.example @@ -8,6 +8,14 @@ region = "us-central1" tenant = "acme" env = "stage" +# Deployment mode. For dependencies only on a Shared VPC, set +# create_runtime = false, network_id to the full host-project network ID, and +# create_psa_connection = false after configuring PSA on that network. +# create_runtime = true +# network_id = "" +# create_psa_connection = true +# redis_transit_encryption = true + # Tenant-supplied secrets. Prefer TF_VAR_litellm_master_key / # TF_VAR_litellm_license / TF_VAR_ui_password env vars so the values don't # end up in a committed tfvars file. All three are optional — when diff --git a/terraform/litellm/gcp/examples/default/variables.tf b/terraform/litellm/gcp/examples/default/variables.tf index 56e5ec88ef8..88b57ce27eb 100644 --- a/terraform/litellm/gcp/examples/default/variables.tf +++ b/terraform/litellm/gcp/examples/default/variables.tf @@ -26,6 +26,30 @@ variable "env" { type = string } +variable "create_runtime" { + description = "Create Cloud Run and load balancer resources." + type = bool + default = true +} + +variable "network_id" { + description = "Existing VPC network resource ID. Empty creates a VPC." + type = string + default = "" +} + +variable "create_psa_connection" { + description = "Create Private Services Access resources." + type = bool + default = true +} + +variable "redis_transit_encryption" { + description = "Enable Memorystore transit encryption." + type = bool + default = true +} + # Sensitive — prefer TF_VAR_litellm_master_key / TF_VAR_litellm_license / # TF_VAR_ui_password so values stay out of any committed tfvars file. variable "litellm_master_key" { diff --git a/terraform/litellm/gcp/iam.tf b/terraform/litellm/gcp/iam.tf index 09df5e7dff0..509e6d48ffd 100644 --- a/terraform/litellm/gcp/iam.tf +++ b/terraform/litellm/gcp/iam.tf @@ -6,6 +6,15 @@ resource "google_service_account" "runtime" { account_id = "${local.name}-runtime" display_name = "LiteLLM Cloud Run runtime" + + lifecycle { + precondition { + condition = !var.create_runtime || var.billing_metrics_endpoint == "" || ( + var.billing_metrics_client_cert_pem != "" && var.billing_metrics_client_key_pem != "" + ) + error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set and create_runtime is true." + } + } } # UI runtime SA — no role bindings. The UI is static nginx with no DB, @@ -14,6 +23,8 @@ resource "google_service_account" "runtime" { # project's serverless service agent (not this SA), so it doesn't need # artifactregistry.reader either. resource "google_service_account" "ui_runtime" { + count = var.create_runtime ? 1 : 0 + account_id = "${local.name}-ui-runtime" display_name = "LiteLLM Cloud Run UI runtime (no data-plane access)" } diff --git a/terraform/litellm/gcp/load_balancer.tf b/terraform/litellm/gcp/load_balancer.tf index 11f30d0f944..57e8af8210f 100644 --- a/terraform/litellm/gcp/load_balancer.tf +++ b/terraform/litellm/gcp/load_balancer.tf @@ -14,77 +14,93 @@ locals { } resource "google_compute_global_address" "lb" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-lb-ip" labels = local.labels } # Serverless NEGs — one per Cloud Run service. resource "google_compute_region_network_endpoint_group" "gateway" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-gateway-neg" region = var.region network_endpoint_type = "SERVERLESS" cloud_run { - service = google_cloud_run_v2_service.gateway.name + service = google_cloud_run_v2_service.gateway[0].name } } resource "google_compute_region_network_endpoint_group" "backend" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-backend-neg" region = var.region network_endpoint_type = "SERVERLESS" cloud_run { - service = google_cloud_run_v2_service.backend.name + service = google_cloud_run_v2_service.backend[0].name } } resource "google_compute_region_network_endpoint_group" "ui" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-ui-neg" region = var.region network_endpoint_type = "SERVERLESS" cloud_run { - service = google_cloud_run_v2_service.ui.name + service = google_cloud_run_v2_service.ui[0].name } } # Backend services wrap each NEG. resource "google_compute_backend_service" "gateway" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-gateway-bs" protocol = "HTTP" load_balancing_scheme = "EXTERNAL_MANAGED" backend { - group = google_compute_region_network_endpoint_group.gateway.id + group = google_compute_region_network_endpoint_group.gateway[0].id } } resource "google_compute_backend_service" "backend" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-backend-bs" protocol = "HTTP" load_balancing_scheme = "EXTERNAL_MANAGED" backend { - group = google_compute_region_network_endpoint_group.backend.id + group = google_compute_region_network_endpoint_group.backend[0].id } } resource "google_compute_backend_service" "ui" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-ui-bs" protocol = "HTTP" load_balancing_scheme = "EXTERNAL_MANAGED" backend { - group = google_compute_region_network_endpoint_group.ui.id + group = google_compute_region_network_endpoint_group.ui[0].id } } # URL map. Default → backend (management API). Path matchers route the # gateway and UI prefixes elsewhere. resource "google_compute_url_map" "this" { + count = var.create_runtime ? 1 : 0 + name = local.name - default_service = google_compute_backend_service.backend.id + default_service = google_compute_backend_service.backend[0].id host_rule { hosts = ["*"] @@ -93,13 +109,13 @@ resource "google_compute_url_map" "this" { path_matcher { name = "main" - default_service = google_compute_backend_service.backend.id + default_service = google_compute_backend_service.backend[0].id # UI paths (catch them before any /v1/* gateway rules so /favicon.ico # and / take precedence). path_rule { paths = local.ui_path_prefixes - service = google_compute_backend_service.ui.id + service = google_compute_backend_service.ui[0].id } # Gateway path prefixes. GCP URL maps cap a path_rule at 10 path globs, @@ -108,7 +124,7 @@ resource "google_compute_url_map" "this" { for_each = { for idx, chunk in chunklist(local.gateway_path_prefixes, 10) : idx => chunk } content { paths = path_rule.value - service = google_compute_backend_service.gateway.id + service = google_compute_backend_service.gateway[0].id } } } @@ -118,7 +134,7 @@ resource "google_compute_url_map" "this" { # target proxy when TLS is enabled; otherwise the regular path-routing # URL map is attached to the HTTP proxy and everything stays plaintext. resource "google_compute_url_map" "https_redirect" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 name = "${local.name}-redirect" default_url_redirect { @@ -129,8 +145,10 @@ resource "google_compute_url_map" "https_redirect" { } resource "google_compute_target_http_proxy" "this" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-http" - url_map = local.tls_enabled ? google_compute_url_map.https_redirect[0].id : google_compute_url_map.this.id + url_map = local.tls_enabled ? google_compute_url_map.https_redirect[0].id : google_compute_url_map.this[0].id # Default-deny on the HTTP-only path: TLS is the supported posture. # Operators must either supply DNS names or explicitly opt in. @@ -143,12 +161,14 @@ resource "google_compute_target_http_proxy" "this" { } resource "google_compute_global_forwarding_rule" "http" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-http" ip_protocol = "TCP" port_range = "80" load_balancing_scheme = "EXTERNAL_MANAGED" - ip_address = google_compute_global_address.lb.address - target = google_compute_target_http_proxy.this.id + ip_address = google_compute_global_address.lb[0].address + target = google_compute_target_http_proxy.this[0].id labels = local.labels } @@ -161,7 +181,7 @@ resource "google_compute_global_forwarding_rule" "http" { # transitions to ACTIVE. resource "google_compute_managed_ssl_certificate" "this" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 # A managed cert's `domains` is immutable, so changing var.lb_domains # forces replacement, and the cert is referenced by the HTTPS target @@ -181,19 +201,19 @@ resource "google_compute_managed_ssl_certificate" "this" { } resource "google_compute_target_https_proxy" "this" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 name = "${local.name}-https" - url_map = google_compute_url_map.this.id + url_map = google_compute_url_map.this[0].id ssl_certificates = [google_compute_managed_ssl_certificate.this[0].id] } resource "google_compute_global_forwarding_rule" "https" { - count = local.tls_enabled ? 1 : 0 + count = var.create_runtime && local.tls_enabled ? 1 : 0 name = "${local.name}-https" ip_protocol = "TCP" port_range = "443" load_balancing_scheme = "EXTERNAL_MANAGED" - ip_address = google_compute_global_address.lb.address + ip_address = google_compute_global_address.lb[0].address target = google_compute_target_https_proxy.this[0].id labels = local.labels } diff --git a/terraform/litellm/gcp/locals.tf b/terraform/litellm/gcp/locals.tf index 9a817eba605..3861413d496 100644 --- a/terraform/litellm/gcp/locals.tf +++ b/terraform/litellm/gcp/locals.tf @@ -21,6 +21,9 @@ locals { var.labels, ) + create_network = var.network_id == "" + network_id = local.create_network ? google_compute_network.this[0].id : var.network_id + gateway_path_prefixes = [ "/v1/chat/*", "/chat/*", "/v1/completions*", "/completions*", @@ -74,7 +77,7 @@ locals { "/ui/*", ] - proxy_config_enabled = length(keys(var.proxy_config)) > 0 + proxy_config_enabled = var.create_runtime && length(keys(var.proxy_config)) > 0 proxy_config_yaml = local.proxy_config_enabled ? yamlencode(var.proxy_config) : "" proxy_config_mount_path = "/etc/litellm" diff --git a/terraform/litellm/gcp/network.tf b/terraform/litellm/gcp/network.tf index a1ccaed02f9..47c7bf94a2b 100644 --- a/terraform/litellm/gcp/network.tf +++ b/terraform/litellm/gcp/network.tf @@ -1,13 +1,17 @@ resource "google_compute_network" "this" { + count = local.create_network ? 1 : 0 + name = local.name auto_create_subnetworks = false routing_mode = "REGIONAL" } resource "google_compute_subnetwork" "this" { + count = local.create_network ? 1 : 0 + name = "${local.name}-${var.region}" region = var.region - network = google_compute_network.this.id + network = google_compute_network.this[0].id ip_cidr_range = var.subnet_cidr private_ip_google_access = true } @@ -16,17 +20,21 @@ resource "google_compute_subnetwork" "this" { # managed services peer with the VPC over the connection below using # addresses from this range. resource "google_compute_global_address" "psa" { + count = var.create_psa_connection ? 1 : 0 + name = "${local.name}-psa" purpose = "VPC_PEERING" address_type = "INTERNAL" prefix_length = 16 - network = google_compute_network.this.id + network = local.network_id } resource "google_service_networking_connection" "psa" { - network = google_compute_network.this.id + count = var.create_psa_connection ? 1 : 0 + + network = local.network_id service = "servicenetworking.googleapis.com" - reserved_peering_ranges = [google_compute_global_address.psa.name] + reserved_peering_ranges = [google_compute_global_address.psa[0].name] } # Serverless VPC Access connector — required so Cloud Run can reach @@ -37,9 +45,11 @@ resource "google_service_networking_connection" "psa" { # for low-to-moderate Cloud Run egress; bump max if your services push # heavy private-network traffic. resource "google_vpc_access_connector" "this" { + count = var.create_runtime ? 1 : 0 + name = "${local.name}-conn" region = var.region - network = google_compute_network.this.name + network = local.network_id ip_cidr_range = var.vpc_connector_cidr min_instances = 2 max_instances = 3 diff --git a/terraform/litellm/gcp/outputs.tf b/terraform/litellm/gcp/outputs.tf index 6f1f1d5ccf4..2a4742f42cf 100644 --- a/terraform/litellm/gcp/outputs.tf +++ b/terraform/litellm/gcp/outputs.tf @@ -1,26 +1,26 @@ output "lb_ip" { - description = "Global anycast IP of the external HTTPS load balancer." - value = google_compute_global_address.lb.address + description = "Global anycast IP of the external HTTPS load balancer. Null when create_runtime is false." + value = var.create_runtime ? one(google_compute_global_address.lb[*].address) : null } output "lb_url" { - description = "Proxy URL. Switches scheme based on whether lb_domains is set; when TLS is enabled the URL points at the first listed domain (since managed certs are tied to the hostname, not the anycast IP). The dashboard is served at /, the API at /v1/*." - value = local.tls_enabled ? "https://${var.lb_domains[0]}" : "http://${google_compute_global_address.lb.address}" + description = "Proxy URL, or null when create_runtime is false. Switches scheme based on whether lb_domains is set." + value = var.create_runtime ? (local.tls_enabled ? "https://${var.lb_domains[0]}" : "http://${one(google_compute_global_address.lb[*].address)}") : null } output "gateway_service_url" { - description = "Default Cloud Run URL for the gateway (bypasses the LB)." - value = google_cloud_run_v2_service.gateway.uri + description = "Default Cloud Run URL for the gateway, or null when create_runtime is false." + value = var.create_runtime ? one(google_cloud_run_v2_service.gateway[*].uri) : null } output "backend_service_url" { - description = "Default Cloud Run URL for the backend (bypasses the LB)." - value = google_cloud_run_v2_service.backend.uri + description = "Default Cloud Run URL for the backend, or null when create_runtime is false." + value = var.create_runtime ? one(google_cloud_run_v2_service.backend[*].uri) : null } output "ui_service_url" { - description = "Default Cloud Run URL for the UI (bypasses the LB)." - value = google_cloud_run_v2_service.ui.uri + description = "Default Cloud Run URL for the UI, or null when create_runtime is false." + value = var.create_runtime ? one(google_cloud_run_v2_service.ui[*].uri) : null } output "cloudsql_writer_ip" { @@ -38,6 +38,36 @@ output "redis_endpoint" { value = "${google_redis_instance.this.host}:${google_redis_instance.this.port}" } +output "runtime_service_account_email" { + description = "Runtime service account email for Cloud Run or GKE Workload Identity." + value = google_service_account.runtime.email +} + +output "redis_host" { + description = "Memorystore Redis host." + value = google_redis_instance.this.host +} + +output "redis_port" { + description = "Memorystore Redis port." + value = google_redis_instance.this.port +} + +output "redis_server_ca_pem" { + description = "Memorystore server CA PEM. Mount it in the pod and set REDIS_SSL=true and REDIS_SSL_CA_CERTS= via extraEnv when transit encryption is enabled." + value = var.redis_transit_encryption ? google_redis_instance.this.server_ca_certs[0].cert : null +} + +output "db_username" { + description = "Cloud SQL application username." + value = var.db_username +} + +output "db_name" { + description = "Cloud SQL database name." + value = var.db_name +} + output "gcs_bucket" { description = "GCS bucket name. Exposed to gateway + backend as GCS_BUCKET_NAME. Reference from proxy_config via `os.environ/GCS_BUCKET_NAME`." value = google_storage_bucket.this.name @@ -54,11 +84,11 @@ output "db_password_secret_id" { } output "migration_run_command" { - description = "Shell command that executes the one-off migration job against Cloud SQL. Run this once after the first apply." - value = format( + description = "Shell command that executes the one-off migration job against Cloud SQL, or null when create_runtime is false." + value = var.create_runtime ? format( "gcloud run jobs execute %s --region %s --project %s --wait", - google_cloud_run_v2_job.migrations.name, + one(google_cloud_run_v2_job.migrations[*].name), var.region, var.project_id, - ) + ) : null } diff --git a/terraform/litellm/gcp/redis.tf b/terraform/litellm/gcp/redis.tf index 0e07c416e85..0602758f090 100644 --- a/terraform/litellm/gcp/redis.tf +++ b/terraform/litellm/gcp/redis.tf @@ -4,7 +4,7 @@ resource "google_redis_instance" "this" { memory_size_gb = var.redis_memory_size_gb region = var.region - authorized_network = google_compute_network.this.id + authorized_network = local.network_id connect_mode = "PRIVATE_SERVICE_ACCESS" redis_version = "REDIS_7_0" @@ -16,7 +16,7 @@ resource "google_redis_instance" "this" { # and passed to the proxy as REDIS_CA_PEM_B64); the proxy decodes it to # /tmp/redis-ca.pem at startup and uses it to validate the rediss:// # handshake. Mirrors `transit_encryption_enabled = true` on AWS. - transit_encryption_mode = "SERVER_AUTHENTICATION" + transit_encryption_mode = var.redis_transit_encryption ? "SERVER_AUTHENTICATION" : "DISABLED" depends_on = [google_service_networking_connection.psa] } diff --git a/terraform/litellm/gcp/tests/deps_only.tftest.hcl b/terraform/litellm/gcp/tests/deps_only.tftest.hcl new file mode 100644 index 00000000000..610c49d5b52 --- /dev/null +++ b/terraform/litellm/gcp/tests/deps_only.tftest.hcl @@ -0,0 +1,175 @@ +mock_provider "google" { + mock_resource "google_redis_instance" { + defaults = { + host = "10.0.0.4" + port = 6379 + server_ca_certs = [{ + cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----" + }] + } + } +} + +mock_provider "google-beta" {} +mock_provider "random" {} + +variables { + project_id = "test-project" + tenant = "tenant" + env = "test" + allow_plaintext_lb = true + image_registry = "us-central1-docker.pkg.dev/test-project/litellm" +} + +run "default_creates_everything" { + command = plan + + assert { + condition = alltrue([ + length(google_compute_network.this) == 1, + length(google_compute_subnetwork.this) == 1, + length(google_compute_global_address.psa) == 1, + length(google_service_networking_connection.psa) == 1, + length(google_vpc_access_connector.this) == 1, + length(google_cloud_run_v2_service.gateway) == 1, + length(google_cloud_run_v2_service.backend) == 1, + length(google_cloud_run_v2_service.ui) == 1, + length(google_cloud_run_v2_job.migrations) == 1, + length(google_compute_global_address.lb) == 1, + length(terraform_data.migration) == 1, + ]) + error_message = "The default mode must create networking, runtime services, the load balancer, and migrations." + } + + assert { + condition = google_redis_instance.this.transit_encryption_mode == "SERVER_AUTHENTICATION" + error_message = "Redis transit encryption must remain enabled by default." + } + + assert { + condition = length(local.shared_env_kv) == 12 + error_message = "The default runtime environment must include GCS and the three Redis TLS entries." + } +} + +run "deps_only_creates_no_runtime" { + command = plan + + variables { + create_runtime = false + proxy_config = { + model_list = [] + } + } + + assert { + condition = alltrue([ + length(google_cloud_run_v2_service.gateway) == 0, + length(google_cloud_run_v2_service.backend) == 0, + length(google_cloud_run_v2_service.ui) == 0, + length(google_cloud_run_v2_job.migrations) == 0, + length(google_cloud_run_v2_service_iam_member.gateway_allusers) == 0, + length(google_cloud_run_v2_service_iam_member.backend_allusers) == 0, + length(google_cloud_run_v2_service_iam_member.ui_allusers) == 0, + length(google_compute_global_address.lb) == 0, + length(google_compute_region_network_endpoint_group.gateway) == 0, + length(google_compute_region_network_endpoint_group.backend) == 0, + length(google_compute_region_network_endpoint_group.ui) == 0, + length(google_compute_backend_service.gateway) == 0, + length(google_compute_backend_service.backend) == 0, + length(google_compute_backend_service.ui) == 0, + length(google_compute_url_map.this) == 0, + length(google_compute_url_map.https_redirect) == 0, + length(google_compute_target_http_proxy.this) == 0, + length(google_compute_global_forwarding_rule.http) == 0, + length(google_compute_managed_ssl_certificate.this) == 0, + length(google_compute_target_https_proxy.this) == 0, + length(google_compute_global_forwarding_rule.https) == 0, + length(terraform_data.migration) == 0, + length(google_vpc_access_connector.this) == 0, + length(google_service_account.ui_runtime) == 0, + length(google_storage_bucket.proxy_config) == 0, + ]) + error_message = "Dependencies-only mode must omit all runtime, load balancer, connector, UI identity, and proxy config resources." + } + + assert { + condition = alltrue([ + google_sql_database_instance.writer.name == "tenant-litellm-test", + google_sql_database_instance.reader.name == "tenant-litellm-test-reader", + google_redis_instance.this.name == "tenant-litellm-test", + google_storage_bucket.this.force_destroy == false, + google_secret_manager_secret.master_key.secret_id == "tenant-litellm-test-master-key", + google_secret_manager_secret.db_password.secret_id == "tenant-litellm-test-db-password", + google_service_account.runtime.account_id == "tenant-litellm-test-runtime", + ]) + error_message = "Dependencies-only mode must retain data stores, secrets, and the runtime service account." + } + + assert { + condition = output.lb_url == null && output.migration_run_command == null + error_message = "Runtime outputs must be null while dependency outputs remain available." + } +} + +run "existing_network_attaches_data_stores" { + command = plan + + variables { + network_id = "projects/host-proj/global/networks/shared" + create_psa_connection = false + create_runtime = false + } + + assert { + condition = alltrue([ + length(google_compute_network.this) == 0, + length(google_compute_subnetwork.this) == 0, + length(google_compute_global_address.psa) == 0, + length(google_service_networking_connection.psa) == 0, + google_sql_database_instance.writer.settings[0].ip_configuration[0].private_network == var.network_id, + google_redis_instance.this.authorized_network == var.network_id, + ]) + error_message = "An existing VPC must receive the Cloud SQL and Memorystore private-network attachments." + } +} + +run "psa_required_without_existing_network" { + command = plan + + variables { + create_psa_connection = false + } + + expect_failures = [ + google_sql_database_instance.writer, + ] +} + +run "redis_plaintext_drops_tls_env" { + command = plan + + variables { + redis_transit_encryption = false + } + + assert { + condition = google_redis_instance.this.transit_encryption_mode == "DISABLED" + error_message = "Redis transit encryption must be disabled when requested." + } + + assert { + condition = length(local.shared_env_kv) == 9 + error_message = "Plaintext Redis mode must include GCS and omit the three Redis TLS entries." + } + + assert { + condition = length([for env in local.shared_env_kv : env if env.name == "REDIS_SSL"]) == 0 + error_message = "Plaintext Redis mode must not set REDIS_SSL." + } + + assert { + condition = length(local.redis_ca_fragment) == 0 + error_message = "Plaintext Redis mode must not decode a Redis CA at startup." + } +} diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index 1162e100bb2..9c68ed3db76 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -79,16 +79,42 @@ variable "ui_password" { sensitive = true } +# ---------- Deployment mode ---------- + +variable "create_runtime" { + description = "Create Cloud Run, load balancer, VPC connector, runtime support resources, and the migration job. Set false for GKE or another external runtime." + type = bool + default = true +} + +variable "network_id" { + description = "Existing VPC network resource ID (`projects//global/networks/`). When set, no VPC or subnet is created. A VPC connector requires this network to be in the deployment project when create_runtime is true." + type = string + default = "" +} + +variable "create_psa_connection" { + description = "Create the Private Services Access range and connection for Cloud SQL and Memorystore. Set false when the existing network already has PSA configured." + type = bool + default = true +} + +variable "redis_transit_encryption" { + description = "Enable Memorystore transit encryption and inject Redis TLS settings into Cloud Run. Set false to use plaintext Redis." + type = bool + default = true +} + # ---------- Networking ---------- variable "subnet_cidr" { - description = "Primary CIDR block for the LiteLLM subnet." + description = "Primary CIDR block for the LiteLLM subnet. Unused when network_id is set." type = string default = "10.40.0.0/16" } variable "vpc_connector_cidr" { - description = "CIDR for the Serverless VPC Access connector. /28 required." + description = "CIDR for the Serverless VPC Access connector. /28 required. Unused when create_runtime is false." type = string default = "10.41.0.0/28" } From 110f654f342897ea438a6c71e91f0078bb4d76fa Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 5 Sep 2026 12:43:02 -0700 Subject: [PATCH 60/61] feat(mcp): renew the stored SSO identity assertion behind ID-JAG (#35401) * feat(mcp): renew the stored SSO identity assertion behind ID-JAG The oauth2_id_jag arm asserts the id_token captured at the user's last interactive SSO login, and nothing ever renewed it, so an agent holding a brokered LiteLLM key could act for that user only until that token's exp. The assertion already carried the IdP refresh token beside it; this redeems it. RefreshingSSOAssertionStore wraps the database reader and satisfies the same protocol, so the egress arm is unchanged. Renewal is lazy and single-flighted per user through the same RefreshCoordinator the authorization_code arm uses, since an IdP that rotates refresh tokens treats two concurrent redemptions as replay. A refusal leaves the expired assertion in place so the reader still challenges the user; an unreachable IdP surfaces as a store outage instead. * fix(mcp): let a cross-replica loser settle the SSO assertion renewal itself Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): satisfy type discipline lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ci): rerun checks after docs main added the missing router setting row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): answer a cross-replica loser retryable instead of re-electing it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): bypass stale assertion cache during renewal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet type-discipline budget after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yassin Kortam Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../per_user_oauth_store.py | 21 +- .../outbound_credentials/resolver.py | 23 +- .../runtime_refresh_coordinator.py | 41 + .../sso_assertion_refresher.py | 469 +++++++++++ .../sso_assertion_store.py | 39 +- .../outbound_credentials/test_resolver.py | 74 ++ .../test_sso_assertion_refresher.py | 794 ++++++++++++++++++ type-discipline-budget.json | 6 +- 8 files changed, 1427 insertions(+), 40 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 9273ddda9cf..5e28396dcfb 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -31,11 +31,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto TokenCacheBackend, TokenStoreUnavailable, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( - RedisDistributedLock, -) -from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( - RedisRefreshCoordinator, +from litellm.proxy._experimental.mcp_server.outbound_credentials.runtime_refresh_coordinator import ( + runtime_refresh_coordinator, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import ( OAuthTokenCacheCodec, @@ -131,23 +128,17 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres ) from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 - redis_cache: Final = user_api_key_cache.redis_cache - if redis_cache is None: + coordinator: Final = runtime_refresh_coordinator() + if coordinator is None: return None, None, False codec: Final = OAuthTokenCacheCodec( encrypt_value_helper, lambda blob: decrypt_value_helper(blob, "mcp_per_user_token", exception_type="debug"), ) - # user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) and the - # Redis client from init_async_client() is partially typed - both are untyped-boundary casts. + # user_api_key_cache satisfies the AsyncCache slice (DualCache types ttl via **kwargs) - an + # untyped-boundary cast. cache: Final[AsyncCache] = user_api_key_cache # pyright: ignore - redis_client: Final = redis_cache.init_async_client() # pyright: ignore - lock: Final = RedisDistributedLock( - redis_client, # pyright: ignore - namespace_key=redis_cache.check_and_fix_namespace, - ) backend: Final = DualCacheTokenCacheBackend(cache, codec) - coordinator: Final = RedisRefreshCoordinator(lock) return backend, coordinator, True diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 404baa14350..8328aae01ab 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -46,11 +46,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( Ok, Result, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + default_sso_assertion_store, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( AssertionStoreUnavailable, - DbSSOAssertionStore, SSOAssertionStore, - SSOIdentityAssertion, + assertion_expired, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.token_endpoint import ( ExchangedToken, @@ -129,7 +131,7 @@ class UpstreamCredentialProvider: self._token_endpoint: TokenEndpointClient = token_endpoint or TokenEndpointClient() self._exchanged_tokens: ExchangedTokenCache = exchanged_tokens or ExchangedTokenCache() self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource() - self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or DbSSOAssertionStore() + self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or default_sso_assertion_store() async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]: match server.config: @@ -246,7 +248,7 @@ class UpstreamCredentialProvider: "Sign in through LiteLLM SSO so the gateway captures one." ) ) - if _assertion_expired(assertion, datetime.now(timezone.utc)): + if assertion_expired(assertion, datetime.now(timezone.utc)): return Error( CredError.of_precondition_required( "The stored IdP identity assertion for this user has expired. Sign in through " @@ -405,19 +407,6 @@ def _id_jag_slot_key(subject: Subject, server: ServerSpec) -> str: return hashlib.sha256(material.encode()).hexdigest() -def _assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: - """Whether the stored assertion's ``exp`` has passed. An assertion carrying no expiry is - treated as usable and left for the IdP to reject, since the store records what the id_token - claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a - stored value that lost its offset compares instead of raising. - """ - expires_at: Final = assertion.expires_at - if expires_at is None: - return False - normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) - return normalized <= now - - def _id_jag_fingerprint(subject_token: str, server_id: str, config: IdJagConfig) -> str: """What the cached leg-2 bearer was minted from: the subject token, the server, and the config. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py new file mode 100644 index 00000000000..e799838b5b7 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -0,0 +1,41 @@ +"""The runtime ``RefreshCoordinator``: cross-replica single-flight when Redis is wired. + +Builds ``RedisRefreshCoordinator`` over the proxy's shared Redis so one refresh runs per key +across the fleet, or returns ``None`` when Redis is absent so the caller keeps the foundation's +in-process default (correct for a single replica). The proxy globals it reads are not ready at +import time, so this is called per composition rather than held as module state. + +Shared by every credential arm that renews a stored grant: a rotating refresh token must be +redeemed once across all workers, so each arm electing its own winner with its own lock shape +would be a bug waiting to differ. +""" + +from __future__ import annotations + +from typing import Final + +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + RefreshCoordinator, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( + RedisDistributedLock, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( + RedisRefreshCoordinator, +) + + +def runtime_refresh_coordinator() -> RefreshCoordinator | None: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # runtime global + + redis_cache: Final = user_api_key_cache.redis_cache + if redis_cache is None: + return None + # The Redis client from init_async_client() is only partially typed; the lock validates every + # reply it depends on, so the untyped boundary is contained here. + redis_client: Final = redis_cache.init_async_client() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm redis wrapper is untyped + lock: Final = RedisDistributedLock( + redis_client, # pyright: ignore[reportArgumentType,reportUnknownArgumentType] # litellm redis wrapper is untyped + namespace_key=redis_cache.check_and_fix_namespace, + ) + return RedisRefreshCoordinator(lock) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py new file mode 100644 index 00000000000..7600fd7ab8a --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py @@ -0,0 +1,469 @@ +"""Renew the stored SSO identity assertion so an ID-JAG agent outlives one id_token. + +The ``oauth2_id_jag`` arm asserts the id_token captured at the user's last interactive sign-in, so +without renewal an agent holding a brokered LiteLLM key can act for that user only until that token's +``exp``, typically an hour, and the sole recovery is another interactive login. The assertion already +carries the IdP refresh token beside it; this module is what redeems it. + +``RefreshingSSOAssertionStore`` wraps any ``SSOAssertionStore`` and satisfies the same protocol, so +the egress arm is unchanged: it still reads one assertion and still judges expiry itself. Renewal is +lazy (only a read that finds a near-expiry assertion triggers one, so IdP traffic tracks actual use, +not the size of the user table) and single-flighted per user through the same ``RefreshCoordinator`` +the ``authorization_code`` arm uses, because an IdP that rotates refresh tokens treats two concurrent +redemptions of one token as replay and can revoke the whole grant chain. + +The refresh is redeemed against the generic-OIDC client the login itself used +(``GENERIC_TOKEN_ENDPOINT`` / ``GENERIC_CLIENT_ID`` / ``GENERIC_CLIENT_SECRET``, which the proxy +reconciles from the stored SSO row into the process environment at startup), authenticated the way +that login authenticated: the non-PKCE path always sends HTTP Basic, while the PKCE path sends the +credentials in the body when ``GENERIC_INCLUDE_CLIENT_ID`` is set, and an IdP application may accept +only one of the two. An assertion can only exist if that client minted it, so no other client could +redeem its refresh token, and no other method is known to be accepted. A deployment whose +``GENERIC_SCOPE`` omits ``offline_access`` captures no refresh token at all, which is why that miss +logs the scope by name rather than failing silently. + +Failures are values internally (``Result[_, RefreshFailure]``). At the store boundary they collapse +onto the protocol's existing two-outcome contract: a refusal returns the expired assertion unchanged +so the reader's own guard challenges the user to sign in again, while a transient IdP failure raises +``AssertionStoreUnavailable`` so the reader answers 503 instead of blaming the user for an outage. + +One ambiguity remains under Redis-coordinated renewal across replicas. A cross-replica loser that +finds the row still expiring after the holder finished cannot tell a refused refresh from a renewal +that could not be recorded. Redeeming itself could consume a refresh token the holder may already +have rotated, so it answers retryable 503 rather than guessing a sign-in challenge. The next +uncontended read settles the outcome itself: a refusal challenges, and a successful refresh persists. +If the holder rotated the token but its write failed, that rotation is lost and the next uncontended +read's refusal challenges, which is the only honest answer because the rotated token was never +recorded. On the refusal path, the loser pays for one retry before that challenge. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Final, Literal, Protocol + +import httpx +from pydantic import SecretStr, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm.exceptions import Timeout +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped +) +from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( + build_token_endpoint_client_auth, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InProcessRefreshCoordinator, + RefreshCoordinator, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.runtime_refresh_coordinator import ( + runtime_refresh_coordinator, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + DbSSOAssertionStore, + SSOAssertionStore, + SSOIdentityAssertion, + assertion_expired, + assertion_from_sso_login, + fetch_sso_identity_assertion, + persist_sso_identity_assertion, +) +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.mcp import MCPTokenEndpointAuthMethod + +_BODY_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(dict[str, object]) + +_REFRESH_GRANT_TYPE: Final = "refresh_token" +# The lock namespace for the one assertion row a user has; the sibling arm keys the same lock by +# server_id, and no server_id can collide with this literal. +_SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion" +# Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of +# the exchange is replaced first. Matches the sibling per-user token store's skew. +_DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0 + + +class AssertionRead(Protocol): + """Reads the user's stored assertion row.""" + + async def __call__(self, user_id: str) -> SSOIdentityAssertion | None: ... + + +class AssertionWrite(Protocol): + """Replaces the user's stored assertion row.""" + + async def __call__(self, user_id: str, assertion: SSOIdentityAssertion) -> None: ... + + +class CoordinatorFactory(Protocol): + """Builds the cross-replica coordinator, or ``None`` when there is no shared lock to build on.""" + + def __call__(self) -> RefreshCoordinator | None: ... + + +class FormPost(Protocol): + """POSTs an OAuth form and hands back the raw response.""" + + async def __call__( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> httpx.Response | None: ... + + +@dataclass(frozen=True, slots=True) +class SSOClientConfig: + """The generic-OIDC client credentials a refresh_token grant has to authenticate as, and how.""" + + token_endpoint: str + client_id: str + client_secret: SecretStr + auth_method: MCPTokenEndpointAuthMethod + + +def sso_client_config(env: Mapping[str, str]) -> SSOClientConfig | None: + """The configured generic-OIDC client, or ``None`` when the deployment has none. + + Read from the process environment because that is where the login path reads it + (``_setup_generic_sso_env_vars``) and where the proxy materializes the stored ``sso_config`` row + at startup, so this resolves to the same client that minted the assertion. ``None`` is an + ordinary state, not an error: a deployment signing in through a provider that captures no + assertion has nothing here to renew, and a client with no secret is not a confidential client + that could redeem one. + + ``auth_method`` is derived from the same ``GENERIC_INCLUDE_CLIENT_ID`` the login reads, because + the two login paths do not agree: the non-PKCE path always authenticates with HTTP Basic, while + the PKCE path puts the credentials in the body when that flag is set. Both capture assertions, so + a constant here would authenticate the renewal differently from the sign-in that produced the + refresh token and 401 against an IdP application registered for only one of the two. + """ + token_endpoint: Final = env.get("GENERIC_TOKEN_ENDPOINT") + client_id: Final = env.get("GENERIC_CLIENT_ID") + client_secret: Final = env.get("GENERIC_CLIENT_SECRET") + if not token_endpoint or not client_id or not client_secret: + return None + includes_client_id: Final = env.get("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" + return SSOClientConfig( + token_endpoint=token_endpoint, + client_id=client_id, + client_secret=SecretStr(client_secret), + auth_method="client_secret_post" if includes_client_id else "client_secret_basic", + ) + + +@dataclass(frozen=True, slots=True) +class RefreshFailure: + """Why a renewal produced nothing, split by what the caller can do about it. + + ``rejected`` is settled: this refresh token will never work again, so the user has to sign in. + ``unavailable`` is transient: the same attempt may succeed in a minute, so telling the user to + sign in again would be a lie about whose problem it is. Both arms carry the same payload, so + this is a ``Literal`` discriminant rather than a ``tagged_union``; consumers still ``match`` on + ``kind`` with an ``assert_never`` tail. + """ + + kind: Literal["rejected", "unavailable"] + detail: str + + @staticmethod + def of_rejected(detail: str) -> RefreshFailure: + return RefreshFailure(kind="rejected", detail=detail) + + @staticmethod + def of_unavailable(detail: str) -> RefreshFailure: + return RefreshFailure(kind="unavailable", detail=detail) + + +class TokenEndpointTransport(Protocol): + """One form POST to the IdP token endpoint, with the refusal/outage split preserved. + + That split is the whole reason this is not the resolver's ``TokenEndpointClient``: that + collaborator maps every non-2xx to ``upstream_unavailable``, which is right for an exchange leg + and wrong here, where a 400 ``invalid_grant`` means the stored refresh token is dead and the user + must act. + """ + + async def post( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> Result[Mapping[str, object], RefreshFailure]: ... + + +async def post_form(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None: + # litellm's httpx handler is only partially typed; nothing but the response object crosses back, + # and the transport below validates its body, so the untyped boundary is contained here. + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped + return await client.post(url, data=form, headers=headers) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportReturnType,reportArgumentType] # litellm http handler is untyped and its stub narrows data=/headers= to dict, which httpx itself does not require + + +class HttpxTokenEndpointTransport: + """The live transport. 4xx is the IdP refusing this grant; anything else is an outage. + + The POST itself is injected so that split, which decides whether the user is challenged or told + to wait, is testable without a live IdP. + """ + + def __init__(self, post: FormPost = post_form) -> None: + self._post = post + + async def post( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> Result[Mapping[str, object], RefreshFailure]: + try: + response: Final = await self._post(url, form, headers) + if response is None: + return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned no response")) + response.raise_for_status() + body: Final = _BODY_ADAPTER.validate_python(response.json()) # pyright: ignore[reportAny] # untyped JSON; the adapter is the type gate + except httpx.HTTPStatusError as exc: + status: Final = exc.response.status_code + if 400 <= status < 500: + return Error(RefreshFailure.of_rejected(f"the IdP refused the refresh with status {status}")) + return Error(RefreshFailure.of_unavailable(f"the IdP token endpoint answered with status {status}")) + except (httpx.RequestError, Timeout) as exc: + return Error(RefreshFailure.of_unavailable(f"the IdP token endpoint is unreachable ({type(exc).__name__})")) + except json.JSONDecodeError: + return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned a non-JSON response")) + except ValidationError: + return Error(RefreshFailure.of_unavailable("the IdP token endpoint returned a non-object response")) + return Ok(body) + + +class SSOAssertionRefresher: + """Redeems the stored refresh token for a current id_token and writes the rotation back. + + Collaborators are injected so the orchestration, the untyped response parsing and the + write-back race are all testable without an IdP or a database. + """ + + def __init__( + self, + transport: TokenEndpointTransport, + *, + client_config: Callable[[], SSOClientConfig | None] = lambda: sso_client_config(os.environ), + read: AssertionRead = fetch_sso_identity_assertion, + write: AssertionWrite = persist_sso_identity_assertion, + ) -> None: + self._transport = transport + self._client_config = client_config + self._read = read + self._write = write + + async def refresh( + self, user_id: str, assertion: SSOIdentityAssertion + ) -> Result[SSOIdentityAssertion, RefreshFailure]: + if assertion.refresh_token is None: + verbose_proxy_logger.warning( + "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and no refresh token was " + "captured with it, so it cannot be renewed without another interactive sign-in. Add " + "'offline_access' to GENERIC_SCOPE so the SSO login captures one.", + user_id, + ) + return Error(RefreshFailure.of_rejected("no refresh token was captured at sign-in")) + config: Final = self._client_config() + if config is None: + verbose_proxy_logger.warning( + "ID-JAG: the stored IdP identity assertion for user_id=%s has expired and cannot be renewed " + "because the generic SSO client is not configured (GENERIC_TOKEN_ENDPOINT, GENERIC_CLIENT_ID, " + "GENERIC_CLIENT_SECRET).", + user_id, + ) + return Error(RefreshFailure.of_rejected("the generic SSO client is not configured")) + + carried_refresh_token: Final = assertion.refresh_token.get_secret_value() + # Whichever method the SSO login used for this client, since that is the one the IdP + # application is known to accept: an assertion only exists to renew because a sign-in already + # authenticated this client that way. + client_auth: Final = build_token_endpoint_client_auth( + auth_method=config.auth_method, + client_id=config.client_id, + client_secret=config.client_secret.get_secret_value(), + ) + form: Final = { # mutable-ok: the RFC 6749 form body is a wire format the HTTP client takes as a mapping + "grant_type": _REFRESH_GRANT_TYPE, + "refresh_token": carried_refresh_token, + **client_auth.body, + } + match await self._transport.post(config.token_endpoint, form, client_auth.headers): + case Error(failure): + return Error(failure) + case Ok(body): + return await self._renewed_from(user_id, assertion, body, carried_refresh_token) + + async def _renewed_from( + self, + user_id: str, + previous: SSOIdentityAssertion, + body: Mapping[str, object], + carried_refresh_token: str, + ) -> Result[SSOIdentityAssertion, RefreshFailure]: + """The renewed assertion, built by the same validator the login path uses. + + A rotated refresh token replaces the stored one; an omitted one carries forward, since an + IdP that does not rotate expects the original to keep working. + """ + rotated: Final = body.get("refresh_token") + renewed: Final = assertion_from_sso_login( + body.get("id_token"), + rotated if isinstance(rotated, str) and rotated else carried_refresh_token, + ) + if renewed is None: + verbose_proxy_logger.warning( + "ID-JAG: the IdP accepted the refresh for user_id=%s but returned no usable id_token, so there " + "is nothing to assert upstream. The SSO client's grant needs the 'openid' scope for the token " + "endpoint to return one on a refresh.", + user_id, + ) + return Error(RefreshFailure.of_rejected("the IdP's refresh response carried no usable id_token")) + failure: Final = await self._store_renewal(user_id, previous, renewed) + if failure is not None: + return Error(failure) + return Ok(renewed) + + async def _store_renewal( + self, user_id: str, previous: SSOIdentityAssertion, renewed: SSOIdentityAssertion + ) -> RefreshFailure | None: + """Write the renewal back, unless the row moved on while this renewal was in flight. + + The row is one per user and last-write-wins, so an interactive sign-in landing mid-renewal + would otherwise be overwritten with a refresh token the IdP has already rotated away, costing + that user a sign-in later. Comparing against the id_token this renewal started from is what + detects that; skipping is safe because the newer row is the one the reader wants anyway. + + A failed write is transient, not settled. The store, not this return value, is what every + caller reads, so a renewal that could not be recorded is a renewal nobody will see; saying so + keeps a database problem answering 503 rather than telling the user to sign in again over it. + """ + try: + current: Final = await self._read(user_id) + if current is not None and current.id_token.get_secret_value() != previous.id_token.get_secret_value(): + verbose_proxy_logger.info( + "ID-JAG: a newer IdP identity assertion for user_id=%s was stored while this renewal was in " + "flight; keeping the stored one.", + user_id, + ) + return None + await self._write(user_id, renewed) + except Exception as exc: # noqa: BLE001 # any storage failure is transient here, never the user's fault + verbose_proxy_logger.warning( + "ID-JAG: could not persist the renewed IdP identity assertion for user_id=%s, so the rotated " + "refresh token is lost and this user will have to sign in again once the renewed token expires: %s", + user_id, + exc, + ) + return RefreshFailure.of_unavailable("the renewed IdP identity assertion could not be persisted") + return None + + +class RefreshingSSOAssertionStore: + """An ``SSOAssertionStore`` that renews a near-expiry assertion before handing it back. + + Reads the inner store; an assertion still comfortably inside its lifetime is returned untouched, + so the common path costs exactly what it did before. Otherwise one renewal runs per user through + the injected ``RefreshCoordinator`` and every caller then re-reads the inner store, which is the + authority: the winner's write is what they all observe, and a renewal the write-back guard + skipped yields the newer assertion that displaced it rather than a private copy. + + A refusal leaves the expired assertion in place for the reader's own guard to reject, so the user + sees the same sign-in-again challenge as before this store existed. A transient IdP failure + raises ``AssertionStoreUnavailable``, the protocol's existing signal for "this is not the user's + fault"; concurrent in-process callers share that outcome, while a cross-replica loser answers 503 + when its re-read still finds the row expiring. On the refusal path that costs the loser one retry, + which then challenges. If the holder rotated the token but its write failed, the rotation is lost + and the next uncontended read's refusal challenges, the only honest answer because that token was + never recorded. + """ + + def __init__( + self, + inner: SSOAssertionStore, + refresher: SSOAssertionRefresher, + *, + fresh_read: AssertionRead, + coordinator_factory: CoordinatorFactory = runtime_refresh_coordinator, + expiry_skew_seconds: float = _DEFAULT_EXPIRY_SKEW_SECONDS, + clock: Callable[[], datetime] = lambda: datetime.now(timezone.utc), + ) -> None: + self._inner = inner + self._refresher = refresher + self._fresh_read = fresh_read + self._coordinator_factory = coordinator_factory + self._in_process_coordinator = InProcessRefreshCoordinator() + self._distributed_coordinator: RefreshCoordinator | None = None + self._skew = timedelta(seconds=expiry_skew_seconds) + self._clock = clock + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + assertion: Final = await self._inner.fetch(user_id) + if not self._expiring(assertion): + return assertion + await self._coordinator().run( + user_id, + _SINGLE_FLIGHT_KEY, + refresh=lambda: self._renew(user_id), + reread=lambda: self._reread_renewed(user_id), + ) + return await self._fresh_read(user_id) + + def _expiring(self, assertion: SSOIdentityAssertion | None) -> bool: + return assertion is not None and assertion_expired(assertion, self._clock() + self._skew) + + def _coordinator(self) -> RefreshCoordinator: + """The cross-replica coordinator once Redis is reachable, else the in-process one. + + Built on first use and kept, because the proxy's Redis client is not wired at import time; + retried while it is absent so a proxy that gains Redis later stops electing per-worker. + """ + if self._distributed_coordinator is None: + self._distributed_coordinator = self._coordinator_factory() + return self._distributed_coordinator or self._in_process_coordinator + + async def _renew(self, user_id: str) -> None: + """The elected renewal, judged from a fresh read so a rotation another replica just landed is + never redeemed again. Returns nothing: the inner store, not this return value, is what every + caller reads afterwards, so the winner and the losers cannot disagree.""" + latest: Final = await self._fresh_read(user_id) + if latest is None or not self._expiring(latest): + return + match await self._refresher.refresh(user_id, latest): + case Ok(_): + return + case Error(failure): + match failure.kind: + case "rejected": + return + case "unavailable": + raise AssertionStoreUnavailable(failure.detail) + assert_never(failure.kind) + + async def _reread_renewed(self, user_id: str) -> None: + """A loser cannot distinguish refusal from an unrecorded renewal without risking token replay. + + It answers retryable 503 instead of guessing a sign-in challenge; the retry runs uncontended + and settles the outcome itself. + """ + latest: Final = await self._fresh_read(user_id) + if self._expiring(latest): + raise AssertionStoreUnavailable( + f"the IdP identity assertion for user_id={user_id} was being renewed by another replica " + "and is not yet current; retry shortly" + ) + + +def default_sso_assertion_store() -> SSOAssertionStore: + """The live read seam for the ``id_jag`` arm: the stored assertion, renewed when it is stale.""" + db_store: Final = DbSSOAssertionStore() + fresh_read: Final = db_store.fetch_uncached + return RefreshingSSOAssertionStore( + db_store, + SSOAssertionRefresher(HttpxTokenEndpointTransport(), read=fresh_read), + fresh_read=fresh_read, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 6552008ca54..f7b92df5ba3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py @@ -127,6 +127,23 @@ def assertion_from_sso_login(id_token: object, refresh_token: object) -> SSOIden ) +def assertion_expired(assertion: SSOIdentityAssertion, now: datetime) -> bool: + """Whether the assertion's ``exp`` has passed at ``now``. An assertion carrying no expiry is + treated as usable and left for the IdP to reject, since the store records what the id_token + claimed rather than imposing a lifetime of its own. A naive ``expires_at`` is read as UTC so a + stored value that lost its offset compares instead of raising. + + Lives beside the model rather than in either reader so the egress guard and the renewal + trigger judge the same field the same way; passing a ``now`` in the future is how a caller + asks "is this about to expire" without a second, driftable predicate. + """ + expires_at: Final = assertion.expires_at + if expires_at is None: + return False + normalized: Final = expires_at if expires_at.tzinfo is not None else expires_at.replace(tzinfo=timezone.utc) + return normalized <= now + + async def ema_assertion_retention_enabled() -> bool: """Whether any MCP server uses ``oauth2_id_jag``, evaluated per login so the gateway only retains bearer material while an EMA upstream exists to spend it on. Judged against the two @@ -146,7 +163,9 @@ async def ema_assertion_retention_enabled() -> bool: return True if prisma_client is None: return False - row = await prisma_client.db.litellm_mcpservertable.find_first(where={"auth_type": MCPAuth.oauth2_id_jag.value}) + row: Final = await prisma_client.db.litellm_mcpservertable.find_first( + where={"auth_type": MCPAuth.oauth2_id_jag.value} + ) return row is not None @@ -158,7 +177,7 @@ async def persist_sso_identity_assertion( if prisma_client is None: return - payload: Final[dict[str, str]] = { + payload: Final = { "id_token": assertion.id_token.get_secret_value(), **({"refresh_token": assertion.refresh_token.get_secret_value()} if assertion.refresh_token else {}), **({"issuer": assertion.issuer} if assertion.issuer else {}), @@ -220,11 +239,13 @@ async def fetch_sso_identity_assertion( class AssertionStoreUnavailable(Exception): - """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down). + """Raised by ``fetch`` when the assertion cannot be read for a transient reason: the DB is + down, or the IdP behind a renewing store could not be reached. Distinct from returning ``None`` for "this user has no captured assertion": an outage must not read as a definite absence, which would tell the user to sign in again over a transient failure, - and it must not escape as an unhandled error on the egress or retry path. Mirrors + and it must not escape as an unhandled error on the egress or retry path. The message names the + real component for the operator log; callers get the reader's generic 503. Mirrors ``TokenStoreUnavailable`` on the sibling per-user OAuth store. """ @@ -257,6 +278,12 @@ class DbSSOAssertionStore: except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence raise AssertionStoreUnavailable(str(exc)) from exc + async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None: + try: + return await _read_assertion_from_db(user_id) + except Exception as exc: # noqa: BLE001 # any driver/storage failure is an outage, not an absence + raise AssertionStoreUnavailable(str(exc)) from exc + async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, new_master_key: str) -> None: """Re-encrypt every stored assertion under ``new_master_key`` during a salt-key rotation, @@ -280,7 +307,9 @@ async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, row.user_id, ) return False - re_encrypted = _STR_ADAPTER.validate_python(encrypt_value_helper(plaintext, new_encryption_key=new_master_key)) + re_encrypted: Final = _STR_ADAPTER.validate_python( + encrypt_value_helper(plaintext, new_encryption_key=new_master_key) + ) await prisma_client.db.litellm_ssoidentityassertion.update( where={"user_id": row.user_id}, data={"assertion_b64": re_encrypted}, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 5e2f2cf97d7..6b3098d9e60 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -9,9 +9,11 @@ returning the stub. import asyncio import logging +import time from datetime import datetime, timedelta, timezone import httpx +import jwt as pyjwt import pytest from pydantic import SecretStr @@ -42,6 +44,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_sto OAuthToken, TokenStoreUnavailable, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + RefreshingSSOAssertionStore, + SSOAssertionRefresher, + SSOClientConfig, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( AssertionStoreUnavailable, SSOIdentityAssertion, @@ -589,6 +596,73 @@ async def test_id_jag_refuses_an_expired_stored_assertion_without_calling_the_id assert endpoint.calls == [] +@pytest.mark.asyncio +async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging(): + """The unattended-agent case end to end: the user last signed in more than an id_token lifetime + ago, so without renewal this is the 412 above. With the renewing store wired the arm resolves, + and leg 1 asserts the renewed token rather than the one that ran out.""" + renewed_id_token = pyjwt.encode( + {"iss": "https://idp.example.com", "sub": "alice", "exp": int(time.time()) + 3600}, + "test-idp-signing-key-32-bytes-long-xxxx", + algorithm="HS256", + ) + expired = SSOIdentityAssertion( + id_token=SecretStr("stale-id-token"), + refresh_token=SecretStr("rt_1"), + expires_at=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + rows = {"alice": expired} + + async def _read(user_id: str) -> SSOIdentityAssertion | None: + return rows.get(user_id) + + async def _write(user_id: str, assertion: SSOIdentityAssertion) -> None: + rows[user_id] = assertion + + class _Inner: + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + return await _read(user_id) + + class _Transport: + async def post(self, url, form, headers): + return Ok({"access_token": "at", "id_token": renewed_id_token}) + + refresher = SSOAssertionRefresher( + _Transport(), + client_config=lambda: SSOClientConfig( + token_endpoint="https://idp.example.com/token", + client_id="litellm", + client_secret=SecretStr("s"), + auth_method="client_secret_basic", + ), + read=_read, + write=_write, + ) + endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) + provider = UpstreamCredentialProvider( + token_endpoint=endpoint, + sso_assertion_store=RefreshingSSOAssertionStore( + _Inner(), refresher, fresh_read=_read, coordinator_factory=lambda: None + ), + ) + + result = await provider.resolve_credentials( + Subject(tenant_id="", subject_id="alice"), _spec(_id_jag_config()) + ) + + assert isinstance(result, Ok) + _, _, leg1_params = endpoint.calls[0] + assert leg1_params["subject_token"] == renewed_id_token + + +def test_the_resolver_defaults_to_the_renewing_assertion_store(): + """A resolver built without collaborators is what production gets, so the default has to renew; + the plain database reader would strand every agent an id_token lifetime after its user's login.""" + provider = UpstreamCredentialProvider() + + assert isinstance(provider._sso_assertion_store, RefreshingSSOAssertionStore) # noqa: SLF001 # the wiring is the assertion + + @pytest.mark.asyncio async def test_id_jag_accepts_a_stored_assertion_that_declares_no_expiry(): endpoint = _FakeTokenEndpoint(_two_leg_ok("final-access")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py new file mode 100644 index 00000000000..d80913f8d33 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py @@ -0,0 +1,794 @@ +"""Tests for renewing the stored SSO identity assertion behind the ID-JAG arm. + +Pins the contract an unattended agent depends on: an assertion that has run out is renewed from the +refresh token captured beside it instead of stranding the agent until its user signs in again, the +IdP sees one redemption per user no matter how many tool calls arrive at once, a rotation is written +back without overwriting a sign-in that landed mid-renewal, and the two failure kinds stay +distinguishable - a dead refresh token still challenges the user, an unreachable IdP does not. +""" + +import asyncio +import base64 +import itertools +import logging +import time +from collections.abc import Awaitable, Callable, Mapping +from datetime import datetime, timedelta, timezone + +import httpx +import jwt as pyjwt +import pytest +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import ( + Error, + Ok, + Result, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_refresher import ( + HttpxTokenEndpointTransport, + RefreshFailure, + RefreshingSSOAssertionStore, + SSOAssertionRefresher, + SSOClientConfig, + default_sso_assertion_store, + sso_client_config, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + AssertionStoreUnavailable, + SSOIdentityAssertion, +) + +SIGNING_KEY = "test-idp-signing-key-32-bytes-long-xxxx" +ISSUER = "https://idp.example.com" +TOKEN_ENDPOINT = "https://idp.example.com/token" + +_CLIENT = SSOClientConfig( + token_endpoint=TOKEN_ENDPOINT, + client_id="litellm", + client_secret=SecretStr("s3cret"), + auth_method="client_secret_basic", +) +_POST_CLIENT = SSOClientConfig( + token_endpoint=TOKEN_ENDPOINT, + client_id="litellm", + client_secret=SecretStr("s3cret"), + auth_method="client_secret_post", +) + + +_MINTED = itertools.count() + + +def _id_token(subject: str = "u1", exp_offset: int = 3600) -> str: + """A distinct token per call. Two mints with the same claims in the same second would encode + identically, which would let a test that means "the renewed token replaced the old one" pass + while comparing a value to itself.""" + return pyjwt.encode( + {"iss": ISSUER, "sub": subject, "exp": int(time.time()) + exp_offset, "jti": f"t{next(_MINTED)}"}, + SIGNING_KEY, + algorithm="HS256", + ) + + +def _stored(id_token: str, *, expires_in: int, refresh_token: str | None = "rt_1") -> SSOIdentityAssertion: + """A row as the SSO callback wrote it: ``expires_in`` seconds from now, mirroring the id_token.""" + return SSOIdentityAssertion( + id_token=SecretStr(id_token), + refresh_token=SecretStr(refresh_token) if refresh_token else None, + issuer=ISSUER, + expires_at=datetime.now(timezone.utc) + timedelta(seconds=expires_in), + ) + + +class _FakeRows: + """The one assertion row per user: the inner read seam and the refresher's read/write pair.""" + + def __init__(self, rows: dict[str, SSOIdentityAssertion] | None = None) -> None: + self.rows: dict[str, SSOIdentityAssertion] = dict(rows or {}) + self.cached_rows: dict[str, SSOIdentityAssertion] = {} + self.reads: list[str] = [] + self.writes: list[tuple[str, SSOIdentityAssertion]] = [] + + async def fetch(self, user_id: str) -> SSOIdentityAssertion | None: + self.reads.append(user_id) + # A real suspension point, so concurrent callers interleave here instead of running to + # completion one at a time and never actually racing. + await asyncio.sleep(0) + return self.cached_rows.get(user_id, self.rows.get(user_id)) + + async def fetch_fresh(self, user_id: str) -> SSOIdentityAssertion | None: + self.reads.append(user_id) + await asyncio.sleep(0) + return self.rows.get(user_id) + + async def write(self, user_id: str, assertion: SSOIdentityAssertion) -> None: + self.writes.append((user_id, assertion)) + self.rows[user_id] = assertion + + +class _FakeTransport: + """Answers every refresh with the same canned result, optionally holding until ``gate`` opens.""" + + def __init__( + self, + response: Result[Mapping[str, object], RefreshFailure], + *, + gate: asyncio.Event | None = None, + on_call: Callable[[], None] | None = None, + ) -> None: + self._response = response + self._gate = gate + self._on_call = on_call + self.calls: list[tuple[str, dict[str, str]]] = [] + self.headers: list[dict[str, str]] = [] + + async def post( + self, url: str, form: Mapping[str, str], headers: Mapping[str, str] + ) -> Result[Mapping[str, object], RefreshFailure]: + self.calls.append((url, dict(form))) + self.headers.append(dict(headers)) + if self._on_call is not None: + self._on_call() + if self._gate is not None: + await self._gate.wait() + return self._response + + +def _renewal(id_token: str, refresh_token: str | None = None) -> Result[Mapping[str, object], RefreshFailure]: + body: dict[str, object] = {"access_token": "at", "id_token": id_token, "token_type": "Bearer"} + return Ok({**body, "refresh_token": refresh_token} if refresh_token else body) + + +def _store( + rows: _FakeRows, + transport: _FakeTransport, + *, + client_config: Callable[[], SSOClientConfig | None] = lambda: _CLIENT, + coordinator_factory: Callable[[], object] = lambda: None, +) -> RefreshingSSOAssertionStore: + refresher = SSOAssertionRefresher(transport, client_config=client_config, read=rows.fetch, write=rows.write) + return RefreshingSSOAssertionStore( + rows, + refresher, + fresh_read=rows.fetch_fresh, + coordinator_factory=coordinator_factory, # pyright: ignore[reportArgumentType] # test doubles stand in for the runtime factory + ) + + +async def _until(predicate: Callable[[], bool]) -> None: + for _ in range(2000): + if predicate(): + return + await asyncio.sleep(0) + raise AssertionError("condition never became true") + + +@pytest.mark.asyncio +async def test_an_expiring_assertion_is_renewed_and_the_renewal_is_what_the_reader_gets(): + """The whole point: an agent calling after its user's id_token ran out keeps working.""" + stale, fresh = _id_token(exp_offset=-1), _id_token() + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(_renewal(fresh)) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == fresh + assert len(transport.calls) == 1 + url, form = transport.calls[0] + assert url == TOKEN_ENDPOINT + assert form["grant_type"] == "refresh_token" + assert form["refresh_token"] == "rt_1" + + +@pytest.mark.asyncio +async def test_a_basic_auth_login_gets_a_basic_auth_refresh(): + """The non-PKCE login always sends HTTP Basic, so the renewal must too; credentials in the body + would 401 against an IdP application registered for Basic.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + await _store(rows, transport).fetch("alice") + + expected = base64.b64encode(b"litellm:s3cret").decode() + assert transport.headers[0]["Authorization"] == f"Basic {expected}" + _url, form = transport.calls[0] + assert "client_secret" not in form + assert "client_id" not in form + + +@pytest.mark.asyncio +async def test_a_body_credential_login_gets_a_body_credential_refresh(): + """The mirror case. A PKCE deployment with GENERIC_INCLUDE_CLIENT_ID set signs in with the + credentials in the body, so Basic here would 401 against an application registered for post; the + renewal has to follow the login rather than a constant.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + await _store(rows, transport, client_config=lambda: _POST_CLIENT).fetch("alice") + + assert "Authorization" not in transport.headers[0] + _url, form = transport.calls[0] + assert form["client_id"] == "litellm" + assert form["client_secret"] == "s3cret" + + +@pytest.mark.parametrize( + ("include_client_id", "expected"), + [ + (None, "client_secret_basic"), + ("false", "client_secret_basic"), + ("TRUE", "client_secret_post"), + ("true", "client_secret_post"), + ], +) +def test_the_auth_method_follows_the_flag_the_login_reads(include_client_id, expected): + """``GENERIC_INCLUDE_CLIENT_ID`` is what the PKCE login branches on, parsed the same way it + parses it, so the renewal cannot pick a method the sign-in did not use.""" + env = { + "GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, + "GENERIC_CLIENT_ID": "litellm", + "GENERIC_CLIENT_SECRET": "s3cret", + **({"GENERIC_INCLUDE_CLIENT_ID": include_client_id} if include_client_id is not None else {}), + } + + config = sso_client_config(env) + + assert config is not None + assert config.auth_method == expected + + +@pytest.mark.asyncio +async def test_an_assertion_well_inside_its_lifetime_never_reaches_the_idp(): + """The common path must cost exactly what it did before this store existed.""" + current = _id_token() + rows = _FakeRows({"alice": _stored(current, expires_in=1800)}) + transport = _FakeTransport(_renewal(_id_token())) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == current + assert transport.calls == [] + assert rows.writes == [] + + +@pytest.mark.asyncio +async def test_renewal_starts_inside_the_skew_rather_than_after_expiry(): + """A token that would die between resolution and the second exchange leg is replaced first.""" + about_to_expire, fresh = _id_token(), _id_token() + assert about_to_expire != fresh + rows = _FakeRows({"alice": _stored(about_to_expire, expires_in=30)}) + transport = _FakeTransport(_renewal(fresh)) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == fresh + + +@pytest.mark.asyncio +async def test_a_user_with_no_stored_assertion_is_still_absent(): + rows = _FakeRows() + transport = _FakeTransport(_renewal(_id_token())) + + assert await _store(rows, transport).fetch("nobody") is None + assert transport.calls == [] + + +@pytest.mark.asyncio +async def test_a_refused_refresh_leaves_the_expired_assertion_for_the_reader_to_reject(): + """A dead refresh token is the user's problem, and the reader's expiry guard is what tells them; + swapping in a renewed-looking value or hiding the row would break that challenge.""" + stale = _id_token(exp_offset=-1) + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(Error(RefreshFailure.of_rejected("the IdP refused the refresh with status 400"))) + + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == stale + assert rows.writes == [] + + +@pytest.mark.asyncio +async def test_an_unreachable_idp_is_a_store_outage_not_a_sign_in_again_challenge(): + """503, not 412: the user has nothing to fix by signing in again while the IdP is down.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(Error(RefreshFailure.of_unavailable("the IdP token endpoint is unreachable"))) + + with pytest.raises(AssertionStoreUnavailable): + await _store(rows, transport).fetch("alice") + + +@pytest.mark.asyncio +async def test_a_missing_refresh_token_names_the_scope_the_operator_has_to_set(caplog): + """Nothing to redeem is the default state of a deployment, so the log has to say what to change + or the feature stays silently inert.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1, refresh_token=None)}) + transport = _FakeTransport(_renewal(_id_token())) + + with caplog.at_level(logging.WARNING): + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert transport.calls == [] + assert "GENERIC_SCOPE" in caplog.text + assert "offline_access" in caplog.text + + +@pytest.mark.asyncio +async def test_an_unconfigured_sso_client_never_calls_the_idp(caplog): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + with caplog.at_level(logging.WARNING): + served = await _store(rows, transport, client_config=lambda: None).fetch("alice") + + assert served is not None + assert transport.calls == [] + assert "GENERIC_TOKEN_ENDPOINT" in caplog.text + + +@pytest.mark.asyncio +async def test_a_refresh_response_carrying_no_id_token_is_refused(caplog): + """An access token is not an identity assertion, so there is nothing to assert upstream.""" + stale = _id_token(exp_offset=-1) + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(Ok({"access_token": "at", "token_type": "Bearer"})) + + with caplog.at_level(logging.WARNING): + served = await _store(rows, transport).fetch("alice") + + assert served is not None + assert served.id_token.get_secret_value() == stale + assert rows.writes == [] + assert "openid" in caplog.text + + +@pytest.mark.asyncio +async def test_a_rotated_refresh_token_replaces_the_stored_one(): + """An IdP that rotates invalidates the old token, so keeping it would cost a sign-in next time.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(), refresh_token="rt_2")) + + await _store(rows, transport).fetch("alice") + + stored = rows.rows["alice"] + assert stored.refresh_token is not None + assert stored.refresh_token.get_secret_value() == "rt_2" + + +@pytest.mark.asyncio +async def test_an_omitted_refresh_token_carries_the_previous_one_forward(): + """An IdP that does not rotate expects the original to keep working; dropping it would strand + the user after exactly one renewal.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + + await _store(rows, transport).fetch("alice") + + stored = rows.rows["alice"] + assert stored.refresh_token is not None + assert stored.refresh_token.get_secret_value() == "rt_1" + + +@pytest.mark.asyncio +async def test_the_renewed_expiry_moves_forward_so_the_next_read_does_not_refresh_again(): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(exp_offset=3600))) + store = _store(rows, transport) + + await store.fetch("alice") + await store.fetch("alice") + + assert len(transport.calls) == 1 + + +async def _explode(user_id: str, assertion: SSOIdentityAssertion) -> None: + raise RuntimeError("write failed") + + +@pytest.mark.asyncio +async def test_a_renewal_that_cannot_be_recorded_is_reported_as_transient(): + """The store is what every caller reads, so a renewal nobody can see is not a success. Calling it + one would hand back a token the gateway failed to record.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + refresher = SSOAssertionRefresher( + _FakeTransport(_renewal(_id_token())), client_config=lambda: _CLIENT, read=rows.fetch, write=_explode + ) + + outcome = await refresher.refresh("alice", rows.rows["alice"]) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_failed_write_does_not_tell_the_user_to_sign_in_again(): + """A database that cannot take the write is not something signing in again fixes, so the reader + has to see an outage rather than the stale row's expiry.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode) + store = RefreshingSSOAssertionStore( + rows, + refresher, + fresh_read=rows.fetch_fresh, + coordinator_factory=lambda: None, # pyright: ignore[reportArgumentType] # test double stands in for the runtime factory + ) + + with pytest.raises(AssertionStoreUnavailable): + await store.fetch("alice") + + assert len(transport.calls) == 1 + + +@pytest.mark.asyncio +async def test_concurrent_reads_for_one_user_redeem_the_refresh_token_once(): + """A burst of tool calls must not replay one refresh token N times: an IdP that rotates reads + that as reuse and can revoke the whole grant chain.""" + gate = asyncio.Event() + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + fresh = _id_token() + transport = _FakeTransport(_renewal(fresh), gate=gate) + store = _store(rows, transport) + + callers = [asyncio.create_task(store.fetch("alice")) for _ in range(8)] + await _until(lambda: len(transport.calls) >= 1 and len(rows.reads) >= 8) + # Guards against a vacuous pass: every caller must have read the expired row and entered the + # renewal branch while the winner is still blocked, otherwise they never raced at all. + assert len(rows.reads) >= 8 + assert not any(task.done() for task in callers) + + gate.set() + served = await asyncio.gather(*callers) + + assert len(transport.calls) == 1 + assert {assertion.id_token.get_secret_value() for assertion in served if assertion is not None} == {fresh} + + +@pytest.mark.asyncio +async def test_concurrent_reads_for_different_users_each_get_their_own_refresh(): + """Single-flight is per user; collapsing across users would leave everyone but one stranded.""" + gate = asyncio.Event() + rows = _FakeRows( + { + "alice": _stored(_id_token("alice", exp_offset=-1), expires_in=-1), + "bob": _stored(_id_token("bob", exp_offset=-1), expires_in=-1), + } + ) + transport = _FakeTransport(_renewal(_id_token()), gate=gate) + store = _store(rows, transport) + + callers = [asyncio.create_task(store.fetch(user)) for user in ("alice", "bob")] + await _until(lambda: len(transport.calls) >= 2) + gate.set() + await asyncio.gather(*callers) + + assert len(transport.calls) == 2 + assert {form["refresh_token"] for _url, form in transport.calls} == {"rt_1"} + + +@pytest.mark.asyncio +async def test_a_renewal_writes_back_when_the_row_did_not_move(): + """The refresh-then-sign-in ordering: nothing displaced the row, so the rotation must land.""" + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + fresh = _id_token() + transport = _FakeTransport(_renewal(fresh, refresh_token="rt_2")) + + served = await _store(rows, transport).fetch("alice") + + assert [user_id for user_id, _assertion in rows.writes] == ["alice"] + assert rows.rows["alice"].id_token.get_secret_value() == fresh + assert served is not None + assert served.id_token.get_secret_value() == fresh + + +@pytest.mark.asyncio +async def test_a_sign_in_landing_mid_renewal_is_not_overwritten(): + """The sign-in-then-refresh ordering. The login wrote a newer assertion while the IdP call was in + flight; overwriting it would put back a refresh token the IdP has already rotated away, costing + that user a sign-in later.""" + from_login = _stored(_id_token("alice"), expires_in=3600, refresh_token="rt_from_login") + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + + def _login_lands() -> None: + rows.rows["alice"] = from_login + + transport = _FakeTransport(_renewal(_id_token(), refresh_token="rt_2"), on_call=_login_lands) + + served = await _store(rows, transport).fetch("alice") + + assert rows.writes == [] + stored = rows.rows["alice"] + assert stored.refresh_token is not None + assert stored.refresh_token.get_secret_value() == "rt_from_login" + assert served is not None + assert served.id_token.get_secret_value() == from_login.id_token.get_secret_value() + + +class _RecordingCoordinator: + """Stands in for the cross-replica coordinator, running the winner's refresh inline.""" + + def __init__(self) -> None: + self.runs: list[tuple[str, str]] = [] + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[None]], + reread: Callable[[], Awaitable[None]], + ) -> None: + self.runs.append((user_id, server_id)) + return await refresh() + + +class _ReplaceThenRefreshCoordinator: + """Replaces the row before running the elected refresh.""" + + def __init__(self, replace: Callable[[], None]) -> None: + self._replace = replace + self.runs: list[tuple[str, str]] = [] + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[None]], + reread: Callable[[], Awaitable[None]], + ) -> None: + self.runs.append((user_id, server_id)) + self._replace() + return await refresh() + + +class _HeldCoordinator: + """Emulates a cross-replica holder finishing before the loser re-reads.""" + + def __init__(self, before_reread: Callable[[], None] | None = None) -> None: + self._before_reread = before_reread + self.runs: list[tuple[str, str]] = [] + + async def run( + self, + user_id: str, + server_id: str, + refresh: Callable[[], Awaitable[None]], + reread: Callable[[], Awaitable[None]], + ) -> None: + self.runs.append((user_id, server_id)) + if self._before_reread is not None: + self._before_reread() + return await reread() + + +@pytest.mark.asyncio +async def test_an_elected_renewal_redeems_the_row_it_re_reads_not_the_one_it_entered_with(): + stale = _id_token(exp_offset=-1) + fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2") + rows = _FakeRows({"alice": _stored(stale, expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + coordinator = _ReplaceThenRefreshCoordinator(lambda: rows.rows.__setitem__("alice", fresh)) + + served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert transport.calls == [] + assert served is not None + assert served.id_token.get_secret_value() == fresh.id_token.get_secret_value() + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_whose_winner_renewed_reads_the_renewal_without_redeeming(): + fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2") + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + coordinator = _HeldCoordinator(before_reread=lambda: rows.rows.__setitem__("alice", fresh)) + + served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert served is fresh + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_rereads_past_a_stale_process_local_cache(): + stale = _stored(_id_token(exp_offset=-1), expires_in=-1) + fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2") + rows = _FakeRows({"alice": fresh}) + rows.cached_rows["alice"] = stale + transport = _FakeTransport(_renewal(_id_token())) + coordinator = _HeldCoordinator() + + served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert served is fresh + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_in_challenge(): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token())) + refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode) + coordinator = _HeldCoordinator() + store = RefreshingSSOAssertionStore( + rows, + refresher, + fresh_read=rows.fetch_fresh, + coordinator_factory=lambda: coordinator, # pyright: ignore[reportArgumentType] # test double stands in for the runtime factory + ) + + with pytest.raises(AssertionStoreUnavailable): + await store.fetch("alice") + + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_a_cross_replica_loser_never_redeems_the_token_the_holder_may_have_rotated(): + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(Error(RefreshFailure.of_rejected("dead"))) + coordinator = _HeldCoordinator() + + with pytest.raises(AssertionStoreUnavailable): + await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice") + + assert transport.calls == [] + assert len(coordinator.runs) == 1 + + +@pytest.mark.asyncio +async def test_the_cross_replica_coordinator_is_used_and_built_once(): + """Redis elects one refresher across the fleet; rebuilding its client per renewal would open a + connection every time.""" + coordinator = _RecordingCoordinator() + builds: list[int] = [] + + def _factory() -> object: + builds.append(1) + return coordinator + + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(exp_offset=-1))) + store = _store(rows, transport, coordinator_factory=_factory) + + await store.fetch("alice") + await store.fetch("alice") + + assert len(builds) == 1 + assert coordinator.runs == [("alice", "sso_identity_assertion"), ("alice", "sso_identity_assertion")] + + +@pytest.mark.asyncio +async def test_the_in_process_coordinator_is_retried_until_redis_appears(): + """A proxy that gains Redis after boot must stop electing a winner per worker.""" + coordinator = _RecordingCoordinator() + available: list[bool] = [False] + + def _factory() -> object | None: + return coordinator if available[0] else None + + rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)}) + transport = _FakeTransport(_renewal(_id_token(exp_offset=-1))) + store = _store(rows, transport, coordinator_factory=_factory) + + await store.fetch("alice") + assert coordinator.runs == [] + + available[0] = True + await store.fetch("alice") + assert coordinator.runs == [("alice", "sso_identity_assertion")] + + +@pytest.mark.parametrize( + "env", + [ + {}, + {"GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"}, + {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_SECRET": "s"}, + {"GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, "GENERIC_CLIENT_ID": "litellm"}, + {"GENERIC_TOKEN_ENDPOINT": "", "GENERIC_CLIENT_ID": "litellm", "GENERIC_CLIENT_SECRET": "s"}, + ], +) +def test_a_partial_sso_client_is_no_client(env): + """Redeeming against a half-configured client would post credentials nowhere useful; the arm + treats it as "cannot renew" and falls back to the sign-in challenge.""" + assert sso_client_config(env) is None + + +def test_the_configured_sso_client_is_the_one_the_login_used(): + config = sso_client_config( + { + "GENERIC_TOKEN_ENDPOINT": TOKEN_ENDPOINT, + "GENERIC_CLIENT_ID": "litellm", + "GENERIC_CLIENT_SECRET": "s3cret", + } + ) + + assert config is not None + assert config.token_endpoint == TOKEN_ENDPOINT + assert config.client_id == "litellm" + assert config.client_secret.get_secret_value() == "s3cret" + + +def test_the_live_store_renews_over_the_database_reader(): + """The composition root has to produce a renewing store, or none of this runs in production.""" + assert isinstance(default_sso_assertion_store(), RefreshingSSOAssertionStore) + + +def _responding(response: httpx.Response | None) -> HttpxTokenEndpointTransport: + async def _post(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None: + return response + + return HttpxTokenEndpointTransport(_post) + + +def _json_response(status: int, payload: dict[str, object]) -> httpx.Response: + return httpx.Response(status, json=payload, request=httpx.Request("POST", TOKEN_ENDPOINT)) + + +@pytest.mark.parametrize("status", [400, 401, 403]) +@pytest.mark.asyncio +async def test_the_idp_declining_the_grant_is_a_refusal_the_user_must_act_on(status): + """A 4xx means this refresh token is finished; calling that an outage would sit the user behind a + 503 forever instead of telling them to sign in.""" + outcome = await _responding(_json_response(status, {"error": "invalid_grant"})).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "rejected" + + +@pytest.mark.parametrize("status", [500, 502, 503]) +@pytest.mark.asyncio +async def test_a_failing_idp_is_an_outage_not_a_refusal(status): + """The refresh token is probably fine; telling the user to sign in again would blame them for + someone else's outage, and would burn their session for nothing.""" + outcome = await _responding(_json_response(status, {})).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_an_unreachable_endpoint_is_an_outage(): + async def _post(url: str, form: Mapping[str, str], headers: Mapping[str, str]) -> httpx.Response | None: + raise httpx.ConnectError("connection refused") + + outcome = await HttpxTokenEndpointTransport(_post).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_non_json_body_is_an_outage(): + response = httpx.Response(200, text="maintenance", request=httpx.Request("POST", TOKEN_ENDPOINT)) + + outcome = await _responding(response).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_missing_response_is_an_outage(): + outcome = await _responding(None).post(TOKEN_ENDPOINT, {}, {}) + + assert isinstance(outcome, Error) + assert outcome.error.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_a_successful_grant_is_handed_back_as_the_parsed_body(): + outcome = await _responding(_json_response(200, {"access_token": "at", "id_token": "idt"})).post( + TOKEN_ENDPOINT, {"grant_type": "refresh_token"}, {} + ) + + assert isinstance(outcome, Ok) + assert outcome.ok["id_token"] == "idt" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 481e3591ce9..225134b4e2b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22181 + "limit": 22180 }, "LIT002": { "limit": 26745 @@ -9,7 +9,7 @@ "limit": 261 }, "LIT004": { - "limit": 40 + "limit": 38 }, "LIT005": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16464 + "limit": 16462 }, "LIT011": { "limit": 5506 From 17e13126cc082134dd2686957c05e74b3e109b05 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 5 Sep 2026 12:43:09 -0700 Subject: [PATCH 61/61] feat(mcp): warn when an oauth2_id_jag server outruns the SSO provider's assertion capture (#35394) * feat(mcp): warn when an oauth2_id_jag server outruns the SSO provider's assertion capture Only the generic OIDC login path captures the IdP id_token that an oauth2_id_jag MCP server spends as its RFC 8693 subject token. Under Google, Microsoft, SAML or no SSO at all, registration succeeds and then every ID-JAG credential resolution fails for every user, with nothing in the logs, the config or the API response to say why. Report the mismatch from the two places it is knowable: when an oauth2_id_jag server is created or updated through the management endpoint, and at SSO callback time when a login hands the arm nothing while such a server is registered. Provider selection mirrors the callback's precedence, so a generic client id sitting behind GOOGLE_CLIENT_ID does not clear the warning. * test(sso): update merged CLI diagnostic patch target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(mcp): warn about the ID-JAG capture gap for config-declared servers and on the SSO debug page (#39350) * feat(sso): surface the ID-JAG capture gap on the SSO debug page /sso/debug/callback is where an operator lands when they are already trying to work out why ID-JAG is failing, so the reason belongs on it. The annotation appears only when the active SSO provider captures no identity assertion AND an oauth2_id_jag server is registered for that gap to break; a deployment without both renders the page it rendered before, byte for byte. Only the provider name and the remedy are rendered, never a configured value, and an unreachable MCP table costs the page its annotation rather than the page itself. The payload carries the one mutable-ok in this work. Conditionally including a member of a JSON document has to construct a mapping, and the rejected alternatives are recorded on the helper so the next reader does not rediscover them. Held out of the diagnosability PR deliberately: that PR is already reviewed and green, and this surface ships with the remaining config-load warning as one follow-up. * feat(mcp): warn at config load when an oauth2_id_jag server outruns the SSO provider's assertion capture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(sso): trim comments on the ID-JAG debug page diagnostic Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(sso): clean up merged imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(sso): satisfy type discipline for diagnostic payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(sso): keep the optional ID-JAG payload member on one line for ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(sso): use Python 3.10-compatible assert_never Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yassin Kortam Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(sso): keep the ID-JAG capture-gap diagnostic out of the unauthenticated debug page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(sso): inject the retention check and log via caplog so the ID-JAG tests pass the test-quality gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(sso): keep the debug-page outage test on the capture-gap path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(sso): annotate the retention check type alias 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> --- .../mcp_server/mcp_server_manager.py | 18 + .../mcp_management_endpoints.py | 22 + .../sso/id_jag_assertion_capture.py | 81 ++++ litellm/proxy/management_endpoints/ui_sso.py | 50 ++- .../mcp_server/test_mcp_server_manager.py | 90 +++++ .../test_id_jag_assertion_capture.py | 117 ++++++ .../test_mcp_management_endpoints.py | 150 +++++++ .../proxy/management_endpoints/test_ui_sso.py | 380 +++++++++++++++++- 8 files changed, 899 insertions(+), 9 deletions(-) create mode 100644 litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bcbcc6bc579..dc1e8db1628 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -159,6 +159,9 @@ from litellm.proxy._types import ( from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + id_jag_assertion_capture_gap_at_startup, +) from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider @@ -1382,6 +1385,20 @@ def _warn_internal_delegate_pkce_if_applicable(server: MCPServer, *, source: str ) +def _warn_config_id_jag_server_outruns_sso(server: MCPServer) -> None: + if server.auth_type != MCPAuth.oauth2_id_jag: + return + gap: Final = id_jag_assertion_capture_gap_at_startup() + if gap is None: + return + verbose_logger.warning( + "MCP server %r (id=%s, source=config) is declared with auth_type=oauth2_id_jag, but %s.", + get_server_prefix(server), + server.server_id, + gap, + ) + + def _deserialize_json_dict(data: str | _StringMap | None) -> dict[str, str] | None: """ Deserialize optional JSON mappings stored in the database. @@ -2393,6 +2410,7 @@ class MCPServerManager: ) self._assign_unique_short_prefix(new_server) _warn_internal_delegate_pkce_if_applicable(new_server, source="config") + _warn_config_id_jag_server_outruns_sso(new_server) self.config_mcp_servers[server_id] = new_server self._set_oauth_discovery_deferred( server_id, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 40cc2e57932..ae266792391 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -64,6 +64,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + id_jag_assertion_capture_gap, +) from litellm.proxy.management_helpers.audit_logs import ( get_audit_log_changed_by, is_audit_logging_enabled, @@ -272,6 +275,22 @@ if MCP_AVAILABLE: _validate_mcp_server_name_fields(payload) _validate_upstream_token_header(payload) + def warn_if_id_jag_server_outruns_sso(server_id: str | None, auth_type: MCPAuth | str | None) -> None: + """Registering an ``oauth2_id_jag`` server under an SSO provider that captures no IdP + identity assertion is a dead configuration: nothing here fails, and then every ID-JAG call + fails for every user with a message that only ever tells them to sign in again. Say it once, + at the moment the admin can still act on it.""" + if auth_type != MCPAuth.oauth2_id_jag: + return + gap = id_jag_assertion_capture_gap() + if gap is None: + return + verbose_proxy_logger.warning( + "MCP server %s is registered with auth_type=oauth2_id_jag, but %s.", + server_id, + gap, + ) + def stamp_omitted_oauth2_flow(payload: NewMCPServerRequest) -> None: """Fallback only: fill in oauth2_flow when an oauth2 create omits it. @@ -1623,6 +1642,8 @@ if MCP_AVAILABLE: detail={"error": f"Error creating mcp server: {e}"}, ) + warn_if_id_jag_server_outruns_sso(new_mcp_server.server_id, new_mcp_server.auth_type) + # Registry refresh is best-effort: the row is already committed, so a # failure here (e.g. an unrelated malformed row in the table) must not # surface as a 500 and orphan the created server, which would push the @@ -2726,6 +2747,7 @@ if MCP_AVAILABLE: status_code=status.HTTP_404_NOT_FOUND, detail={"error": f"MCP Server not found, passed server_id={payload.server_id}"}, ) + warn_if_id_jag_server_outruns_sso(mcp_server_record_updated.server_id, mcp_server_record_updated.auth_type) await global_mcp_server_manager.update_server(mcp_server_record_updated) # Ensure registry is up to date by reloading from database diff --git a/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py new file mode 100644 index 00000000000..404fdfc83a9 --- /dev/null +++ b/litellm/proxy/management_endpoints/sso/id_jag_assertion_capture.py @@ -0,0 +1,81 @@ +"""Whether the SSO provider the login callback dispatches to can capture an IdP identity assertion. + +An ``oauth2_id_jag`` MCP server spends the ``id_token`` captured at SSO login as its RFC 8693 +subject token. Only the generic OIDC login path reaches a token response the gateway retains one +from, so a deployment whose SSO runs through Google, Microsoft or SAML never stores an assertion +and every store-sourced ID-JAG exchange fails for every user, however many times they sign in. +Neither side can see that alone: the MCP registration knows nothing about SSO and the login knows +nothing about MCP. This module is the one shared answer both warn from. +""" + +from __future__ import annotations + +import os +from enum import Enum + +from typing_extensions import assert_never + +from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler + +_GENERIC_OIDC_REMEDY = ( + "Point SSO at the generic OIDC provider (GENERIC_CLIENT_ID), the one login path whose token " + "response the gateway retains an id_token from" +) + + +class ActiveSSOProvider(str, Enum): + google = "google" + microsoft = "microsoft" + generic = "generic" + saml = "saml" + none = "none" + + +def active_sso_provider() -> ActiveSSOProvider: + """The provider the SSO callback will dispatch to. + + Mirrors the callback's precedence rather than reporting everything configured: an environment + carrying both GOOGLE_CLIENT_ID and GENERIC_CLIENT_ID runs the Google branch, so it must report + Google. Presence is judged the way the callback judges it, so a client id set to the empty + string still selects that branch here. + """ + if os.getenv("GOOGLE_CLIENT_ID") is not None: + return ActiveSSOProvider.google + if os.getenv("MICROSOFT_CLIENT_ID") is not None: + return ActiveSSOProvider.microsoft + if os.getenv("GENERIC_CLIENT_ID") is not None: + return ActiveSSOProvider.generic + if SAMLAuthHandler.is_saml_configured(): + return ActiveSSOProvider.saml + return ActiveSSOProvider.none + + +def id_jag_assertion_capture_gap() -> str | None: + """Why ID-JAG cannot work under the active SSO provider, phrased for an operator reading a log, + or ``None`` when that provider does capture an assertion.""" + provider = active_sso_provider() + match provider: + case ActiveSSOProvider.generic: + return None + case ActiveSSOProvider.none: + return ( + "no SSO provider is configured, so no IdP identity assertion is ever captured and " + f"ID-JAG credential resolution fails for every user. {_GENERIC_OIDC_REMEDY}" + ) + case ActiveSSOProvider.google | ActiveSSOProvider.microsoft | ActiveSSOProvider.saml: + return ( + f"the active SSO provider ({provider.value}) has no identity-assertion capture path, so no " + "IdP id_token is ever stored and ID-JAG credential resolution fails for every user no matter " + f"how often they sign in. {_GENERIC_OIDC_REMEDY}" + ) + case _: + assert_never(provider) + + +def id_jag_assertion_capture_gap_at_startup() -> str | None: + """Config load runs before SSO settings stored in the database are reconciled into the process + environment, so an unresolved provider at that point is not yet a gap; the SSO callback reports it + once a login happens.""" + if active_sso_provider() is ActiveSSOProvider.none: + return None + return id_jag_assertion_capture_gap() diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 84150ef7935..3e6434a5afd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,7 +16,7 @@ import json import os import re import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from html import escape from types import MappingProxyType @@ -29,6 +29,7 @@ from typing import ( NoReturn, Optional, Protocol, + TypeAlias, Union, cast, overload, @@ -70,6 +71,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( SSOIdentityAssertion, assertion_from_sso_login, + ema_assertion_retention_enabled, retain_sso_identity_assertion_for_ema, ) from litellm.proxy._types import ( @@ -105,6 +107,9 @@ from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + id_jag_assertion_capture_gap, +) from litellm.proxy.management_endpoints.sso.saml_sso import SAMLAuthHandler from litellm.proxy.management_endpoints.sso_helper_utils import ( check_is_admin_only_access, @@ -1677,6 +1682,46 @@ async def get_generic_sso_response( return result or {}, received_response, access_token_payload, sso_assertion +RetentionCheck: TypeAlias = Callable[[], Awaitable[bool]] # mutable-ok: Callable parameter syntax + + +async def warn_if_id_jag_assertion_uncaptured( + assertion: SSOIdentityAssertion | None, *, retention_enabled: RetentionCheck | None = None +) -> None: + """Say, at the one moment it is knowable, that this login gave an ``oauth2_id_jag`` server + nothing to spend. Without it the operator only ever sees the per-request failure, which cannot + tell a user who has never signed in from a provider that will never capture. Kept strictly + diagnostic: a store outage is swallowed, since a login must not fail over a log line.""" + if assertion is not None: + return + try: + check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled + if not await check(): + return + except Exception as exc: # noqa: BLE001 # diagnostics must never break the login + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers after SSO login: %s", exc) + return + gap: Final = id_jag_assertion_capture_gap() + verbose_proxy_logger.warning( + "SSO login captured no IdP identity assertion while an oauth2_id_jag MCP server is registered: %s", + gap if gap is not None else "the identity provider's token response carried no usable id_token", + ) + + +async def warn_if_id_jag_capture_gap(*, retention_enabled: RetentionCheck | None = None) -> None: + gap: Final = id_jag_assertion_capture_gap() + if gap is None: + return + try: + check: Final = retention_enabled if retention_enabled is not None else ema_assertion_retention_enabled + if not await check(): + return + except Exception as exc: # noqa: BLE001 # diagnostics must never break the page they annotate + verbose_proxy_logger.debug("Could not check for oauth2_id_jag MCP servers: %s", exc) + return + verbose_proxy_logger.warning("SSO debug callback ran with an oauth2_id_jag capture gap: %s", gap) + + async def create_team_member_add_task(team_id, user_info): """Create a task for adding a member to a team.""" try: @@ -2269,6 +2314,7 @@ async def _complete_cli_sso_callback_session( raise HTTPException(status_code=500, detail="Failed to retrieve user information from SSO") await retain_sso_identity_assertion_for_ema(user_id=user_info.user_id, assertion=sso_assertion) + await warn_if_id_jag_assertion_uncaptured(sso_assertion) teams: list[str] = [] if hasattr(user_info, "teams") and user_info.teams: @@ -3599,6 +3645,7 @@ class SSOAuthenticationHandler: if isinstance(user_id, str) and user_id: await retain_sso_identity_assertion_for_ema(user_id=user_id, assertion=sso_assertion) + await warn_if_id_jag_assertion_uncaptured(sso_assertion) disabled_non_admin_personal_key_creation: Final = get_disabled_non_admin_personal_key_creation() litellm_dashboard_ui = get_custom_url(request_base_url=str(request.base_url), route="ui/") @@ -4733,6 +4780,7 @@ async def debug_sso_callback(request: Request): safe_raw_claims: Final = {k: v for k, v in (received_response or {}).items() if k not in _OAUTH_TOKEN_FIELDS} safe_access_token_claims = {k: v for k, v in (access_token_payload or {}).items() if k not in _OAUTH_TOKEN_FIELDS} + await warn_if_id_jag_capture_gap() sso_payload: Final = { "parsed_by_proxy": filtered_result, "raw_claims": safe_raw_claims, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 34dc067e7a3..e3ed48713aa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -461,6 +461,30 @@ class TestMCPServerManager: base.update(overrides) return {"m2mserver": base} + def _id_jag_config(self): + return { + "idjag_server": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2_id_jag, + "client_id": "cid", + "client_secret": "csec", + "token_exchange_endpoint": "https://idp.example.com/token", + "id_jag_resource_token_endpoint": "https://resource.example.com/token", + "id_jag_resource": "https://resource.example.com", + } + } + + def _clear_sso_env(self, monkeypatch): + for env_var in ( + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + ): + monkeypatch.delenv(env_var, raising=False) + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on"]) def test_mcp_oauth_discovery_on_startup_true_values(self, value): with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": value}): @@ -1130,6 +1154,72 @@ class TestMCPServerManager: server = next(iter(manager.config_mcp_servers.values())) assert server.oauth2_flow is None + @pytest.mark.asyncio + async def test_load_servers_from_config_warns_for_id_jag_with_google_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + warnings = [message for message in caplog.messages if "oauth2_id_jag" in message] + assert len(warnings) == 1 + assert "idjag_server" in warnings[0] + assert "GENERIC_CLIENT_ID" in warnings[0] + + @pytest.mark.asyncio + async def test_load_servers_from_config_does_not_warn_for_id_jag_without_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) + + @pytest.mark.asyncio + async def test_load_servers_from_config_does_not_warn_for_api_key_with_google_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + manager = MCPServerManager() + config = { + "api_key_server": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.api_key, + "auth_value": "upstream-secret", + } + } + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(config) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) + + @pytest.mark.asyncio + async def test_load_servers_from_config_does_not_warn_for_id_jag_with_generic_sso(self, monkeypatch, caplog): + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + manager = MCPServerManager() + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + patch.object(manager, "_hydrate_config_servers_dcr_clients", new=AsyncMock()), + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + await manager.load_servers_from_config(self._id_jag_config()) + + assert not any("oauth2_id_jag" in message for message in caplog.messages) + def _client_forwarded_config(self, auth_type, **overrides): base = { "url": "https://example.com/mcp", diff --git a/tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py b/tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py new file mode 100644 index 00000000000..ff4fbbfb695 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_id_jag_assertion_capture.py @@ -0,0 +1,117 @@ +import pytest + +from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import ( + ActiveSSOProvider, + active_sso_provider, + id_jag_assertion_capture_gap, + id_jag_assertion_capture_gap_at_startup, +) + +_SSO_ENV_VARS = ( + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", +) + + +@pytest.fixture(autouse=True) +def _isolated_sso_env(monkeypatch): + """Every SSO selector is read from the process environment, so a value left behind by + another test would silently decide this one's answer.""" + for name in _SSO_ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +class TestActiveSSOProviderMirrorsTheCallback: + """The gap warning is only as good as its agreement with the branch the login callback + actually takes, so provider selection is asserted branch by branch, including the + precedence that makes a co-configured generic client unreachable.""" + + def test_google_client_id_selects_google(self, monkeypatch): + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + assert active_sso_provider() is ActiveSSOProvider.google + + def test_microsoft_client_id_selects_microsoft(self, monkeypatch): + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid") + assert active_sso_provider() is ActiveSSOProvider.microsoft + + def test_generic_client_id_selects_generic(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert active_sso_provider() is ActiveSSOProvider.generic + + def test_saml_metadata_selects_saml(self, monkeypatch): + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert active_sso_provider() is ActiveSSOProvider.saml + + def test_nothing_configured_selects_none(self): + assert active_sso_provider() is ActiveSSOProvider.none + + def test_google_outranks_a_co_configured_generic_client(self, monkeypatch): + """The callback tests GOOGLE_CLIENT_ID first, so the generic arm never runs here and + no assertion is captured; reporting generic would clear a gap that is still open.""" + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert active_sso_provider() is ActiveSSOProvider.google + + def test_microsoft_outranks_a_co_configured_generic_client(self, monkeypatch): + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "ms-cid") + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert active_sso_provider() is ActiveSSOProvider.microsoft + + def test_generic_outranks_saml(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + monkeypatch.setenv("SAML_IDP_METADATA_URL", "https://idp.example.com/metadata") + assert active_sso_provider() is ActiveSSOProvider.generic + + +class TestIdJagAssertionCaptureGap: + def test_generic_oidc_has_no_gap(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert id_jag_assertion_capture_gap() is None + + @pytest.mark.parametrize( + "env_var, provider_label", + [ + ("GOOGLE_CLIENT_ID", "google"), + ("MICROSOFT_CLIENT_ID", "microsoft"), + ("SAML_IDP_METADATA_URL", "saml"), + ], + ) + def test_non_capturing_provider_is_named_with_the_remedy(self, monkeypatch, env_var, provider_label): + monkeypatch.setenv(env_var, "configured") + gap = id_jag_assertion_capture_gap() + assert gap is not None + assert provider_label in gap + assert "GENERIC_CLIENT_ID" in gap + + def test_no_sso_configured_reports_a_gap(self): + gap = id_jag_assertion_capture_gap() + assert gap is not None + assert "no SSO provider is configured" in gap + + def test_google_beside_generic_still_reports_a_gap(self, monkeypatch): + """The precedence trap in operator terms: adding a generic client id without removing + GOOGLE_CLIENT_ID does not fix the deployment, so the gap must not clear.""" + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + gap = id_jag_assertion_capture_gap() + assert gap is not None + assert "google" in gap + + +class TestIdJagAssertionCaptureGapAtStartup: + def test_no_provider_at_startup_is_not_yet_a_gap(self): + assert id_jag_assertion_capture_gap_at_startup() is None + + def test_google_provider_at_startup_reports_the_capture_gap(self, monkeypatch): + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-cid") + startup_gap = id_jag_assertion_capture_gap_at_startup() + callback_gap = id_jag_assertion_capture_gap() + assert startup_gap is not None + assert startup_gap == callback_gap + + def test_generic_provider_at_startup_has_no_gap(self, monkeypatch): + monkeypatch.setenv("GENERIC_CLIENT_ID", "generic-cid") + assert id_jag_assertion_capture_gap_at_startup() is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index adab3538b58..71ff7de89b0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2,6 +2,7 @@ import os import sys import types import json +import logging from contextlib import ExitStack from datetime import datetime, timedelta from types import SimpleNamespace @@ -3840,6 +3841,155 @@ class TestAddMCPServerAtomicity: mock_manager.reload_servers_from_database.assert_not_awaited() +class TestIdJagRegistrationWarnsAboutTheSSOGap: + """An `oauth2_id_jag` server only ever works when the login path captures an IdP identity + assertion, and only the generic OIDC arm does. Registering one under Google or Microsoft + succeeds and then fails for every user on every call, so the mismatch has to be said at + registration time, while the admin is still looking at the configuration.""" + + @staticmethod + def _clear_sso_env(monkeypatch): + for name in ( + "GOOGLE_CLIENT_ID", + "MICROSOFT_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + ): + monkeypatch.delenv(name, raising=False) + + @staticmethod + def _id_jag_warnings(caplog) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "oauth2_id_jag" in record.getMessage() + ] + + @staticmethod + def _server_record(auth_type) -> LiteLLM_MCPServerTable: + record = generate_mock_mcp_server_db_record(server_id="ema-1", alias="ema") + record.auth_type = auth_type + return record + + async def _run_create(self, monkeypatch, provider_env, auth_type, caplog): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + self._clear_sso_env(monkeypatch) + for name, value in provider_env.items(): + monkeypatch.setenv(name, value) + + mock_manager = MagicMock() + mock_manager.add_server = AsyncMock() + mock_manager.reload_servers_from_database = AsyncMock() + + with ( + patch( # test-quality-ok: endpoint test stubs the Prisma client lookup + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint test stubs MCP server creation + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=self._server_record(auth_type)), + ), + patch( # test-quality-ok: endpoint reads the global MCP manager + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await add_mcp_server( + payload=NewMCPServerRequest( + alias="ema", + url="https://ema.example.com/mcp", + transport=MCPTransport.http, + ), + user_api_key_dict=generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ), + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "provider_env, expected_fragment", + [ + ({"GOOGLE_CLIENT_ID": "cid"}, "google"), + ({"MICROSOFT_CLIENT_ID": "cid"}, "microsoft"), + ({"SAML_IDP_METADATA_URL": "https://idp.example.com/metadata"}, "saml"), + ({}, "no SSO provider is configured"), + ], + ) + async def test_create_warns_under_a_provider_that_captures_nothing( + self, monkeypatch, caplog, provider_env, expected_fragment + ): + await self._run_create(monkeypatch, provider_env, MCPAuth.oauth2_id_jag, caplog) + warnings = self._id_jag_warnings(caplog) + assert len(warnings) == 1 + assert expected_fragment in str(warnings[0]) + assert "ema-1" in str(warnings[0]) + + @pytest.mark.asyncio + async def test_create_is_silent_under_generic_oidc(self, monkeypatch, caplog): + await self._run_create(monkeypatch, {"GENERIC_CLIENT_ID": "cid"}, MCPAuth.oauth2_id_jag, caplog) + assert self._id_jag_warnings(caplog) == [] + + @pytest.mark.asyncio + async def test_create_is_silent_for_other_auth_types(self, monkeypatch, caplog): + """Nothing but the id_jag arm sources credentials from a stored SSO assertion, so no + other server registered under Google has anything to warn about.""" + await self._run_create(monkeypatch, {"GOOGLE_CLIENT_ID": "cid"}, MCPAuth.api_key, caplog) + assert self._id_jag_warnings(caplog) == [] + + @pytest.mark.asyncio + async def test_update_to_id_jag_warns(self, monkeypatch, caplog): + """Switching an existing server onto id_jag opens the same gap a create does.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + edit_mcp_server, + ) + + self._clear_sso_env(monkeypatch) + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + + mock_manager = MagicMock() + mock_manager.update_server = AsyncMock() + mock_manager.reload_servers_from_database = AsyncMock() + with ( + patch( # test-quality-ok: endpoint test stubs the Prisma client lookup + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( # test-quality-ok: endpoint test stubs the MCP server lookup + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=self._server_record(MCPAuth.api_key)), + ), + patch( # test-quality-ok: endpoint test stubs MCP server updates + "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", + AsyncMock(return_value=self._server_record(MCPAuth.oauth2_id_jag)), + ), + patch( # test-quality-ok: endpoint test stubs credential cleanup + "litellm.proxy.management_endpoints.mcp_management_endpoints.purge_user_oauth_credentials_for_server", + AsyncMock(return_value=0), + ), + patch( # test-quality-ok: endpoint reads the global MCP manager + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await edit_mcp_server( + payload=UpdateMCPServerRequest(server_id="ema-1", auth_type=MCPAuth.oauth2_id_jag), + user_api_key_dict=generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ), + ) + + warnings = self._id_jag_warnings(caplog) + assert len(warnings) == 1 + assert "google" in str(warnings[0]) + + class TestHealthCheckServers: """Test suite for health check servers endpoint""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 5dfff53f7c3..8d8bc15f9be 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1,17 +1,16 @@ import asyncio import json +import logging import os -from contextlib import asynccontextmanager +from contextlib import ExitStack, asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException, Request -from litellm._uuid import uuid - - import litellm +from litellm._uuid import uuid from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO @@ -1615,8 +1614,8 @@ async def test_get_generic_sso_response_with_empty_headers(): async def test_get_generic_sso_response_includes_token_claims_when_enabled(monkeypatch): import jwt as pyjwt - from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) mock_jwt_handler = MagicMock(spec=JWTHandler) @@ -2321,10 +2320,10 @@ class TestCustomUISSO: async def test_handle_custom_ui_sso_sign_in_success(self): """Test successful custom UI SSO sign-in with valid headers""" from fastapi_sso.sso.base import OpenID - from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler # Mock request with custom headers @@ -2400,6 +2399,7 @@ class TestCustomUISSO: from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler mock_request = MagicMock(spec=Request) @@ -2436,10 +2436,10 @@ class TestCustomUISSO: and its methods are called with the correct parameters """ from fastapi_sso.sso.base import OpenID - from litellm_enterprise.proxy.auth.custom_sso_handler import ( EnterpriseCustomSSOHandler, ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler # Create a real custom handler class instance @@ -8167,6 +8167,128 @@ async def test_debug_sso_callback_handles_missing_raw_response(): assert "user@example.com" in body +# ── The debug page is where an operator lands when ID-JAG is failing ────────── + +_GOOGLE_DEBUG_CLIENT_ID = "debug-google-client-id" +_GENERIC_DEBUG_CLIENT_ID = "debug-generic-client-id" + + +async def _render_debug_page(provider_env, id_jag_registered, force_inert=False): + """Drive /sso/debug/callback and return the raw response body.""" + from litellm.proxy.management_endpoints.ui_sso import GoogleSSOHandler, debug_sso_callback + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://proxy.example.com/" + mock_request.cookies = {} + mock_request.query_params = {} + + parsed = {"sub": "user_123", "email": "u@example.com"} + + async def fake_generic(**kwargs): + return parsed, {"sub": "user_123"}, {"scope": "openid"}, None + + async def fake_google(**kwargs): + return parsed + + stack = [ + patch.dict(os.environ, provider_env, clear=False), + patch( # test-quality-ok: endpoint test stubs the upstream generic IdP boundary + "litellm.proxy.management_endpoints.ui_sso.get_generic_sso_response", side_effect=fake_generic + ), + patch.object( # test-quality-ok: endpoint test stubs the upstream Google IdP boundary + GoogleSSOHandler, "get_google_callback_response", side_effect=fake_google + ), + patch( # test-quality-ok: debug endpoint reads this module global without an injection seam + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=id_jag_registered), + ), + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals + ] + if force_inert: + stack.append( + patch( # test-quality-ok: force-inert reference isolates the endpoint's pre-change response + "litellm.proxy.management_endpoints.ui_sso.warn_if_id_jag_capture_gap", + AsyncMock(return_value=None), + ) + ) + + with ExitStack() as es: + for ctx in stack: + es.enter_context(ctx) + for var in ("MICROSOFT_CLIENT_ID", "GOOGLE_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + if var not in provider_env: + os.environ.pop(var, None) + response = await debug_sso_callback(mock_request) + + return response.body.decode() + + +@pytest.mark.asyncio +async def test_debug_page_logs_the_capture_gap_but_never_renders_it(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + body = await _render_debug_page({"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=True) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "google" in warnings[0] + assert "GENERIC_CLIENT_ID" in warnings[0] + assert "id_jag" not in body + assert "GENERIC_CLIENT_ID" not in body + + +@pytest.mark.asyncio +async def test_debug_page_is_byte_identical_when_the_provider_captures(): + """A deployment with no gap must get the page it got before this change, to the byte. The + comparison is against the endpoint with the diagnostic forced inert, not against a guess.""" + with_feature = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, id_jag_registered=True + ) + pre_change = await _render_debug_page( + {"GENERIC_CLIENT_ID": _GENERIC_DEBUG_CLIENT_ID}, + id_jag_registered=True, + force_inert=True, + ) + + assert with_feature == pre_change + assert "id_jag" not in with_feature + + +@pytest.mark.asyncio +async def test_debug_page_is_byte_identical_when_no_id_jag_server_is_registered(): + """Most deployments run Google SSO and no id_jag server at all; their debug page must not + grow an ID-JAG section about a feature they do not use.""" + with_feature = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, id_jag_registered=False + ) + pre_change = await _render_debug_page( + {"GOOGLE_CLIENT_ID": _GOOGLE_DEBUG_CLIENT_ID}, + id_jag_registered=False, + force_inert=True, + ) + + assert with_feature == pre_change + assert "id_jag" not in with_feature + + +@pytest.mark.asyncio +async def test_debug_page_survives_a_store_outage(monkeypatch, caplog): + """The page's job is to render claims; an unreachable MCP table must cost it the annotation, + not the page.""" + from litellm.proxy.management_endpoints.ui_sso import warn_if_id_jag_capture_gap + + monkeypatch.setenv("GOOGLE_CLIENT_ID", _GOOGLE_DEBUG_CLIENT_ID) + retention_check = AsyncMock(side_effect=Exception("db down")) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + assert await warn_if_id_jag_capture_gap(retention_enabled=retention_check) is None + + retention_check.assert_awaited_once() + + assert _id_jag_gap_warnings(caplog) == [] + + async def _render_legacy_login_page(env_overrides, general_settings): from litellm.proxy.management_endpoints.ui_sso import google_login @@ -8261,8 +8383,8 @@ async def test_saml_callback_enforces_free_sso_user_limit_after_validation(): that /sso/key/generate enforces; the ACS re-checks it after validating the assertion, so the entitlement DB query never runs on unvalidated input.""" from litellm.proxy._types import ProxyException - from litellm.proxy.management_endpoints.ui_sso import saml_callback from litellm.proxy.management_endpoints.types import CustomOpenID + from litellm.proxy.management_endpoints.ui_sso import saml_callback call_order: list[str] = [] @@ -8681,6 +8803,248 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): assert response.status_code == 200 +def _id_jag_gap_warnings(caplog) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "oauth2_id_jag" in record.getMessage() + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_env, expected_fragment", + [ + ({"GOOGLE_CLIENT_ID": "cid"}, "google"), + ({"MICROSOFT_CLIENT_ID": "cid", "MICROSOFT_TENANT": "t"}, "microsoft"), + ({}, "no SSO provider is configured"), + ], +) +async def test_uncaptured_assertion_warns_when_an_id_jag_server_is_registered( + monkeypatch, caplog, provider_env, expected_fragment +): + """A provider with no capture path leaves ID-JAG permanently broken, and the only place + that is knowable is the login itself; without this line the operator sees nothing at all.""" + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "GENERIC_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(name, raising=False) + for name, value in provider_env.items(): + monkeypatch.setenv(name, value) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True)) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert expected_fragment in str(warnings[0]) + + +@pytest.mark.asyncio +async def test_generic_provider_that_returned_no_id_token_still_warns(monkeypatch, caplog): + """Generic OIDC has a capture path, so there is no configuration gap to report; the login + still handed the id_jag arm nothing, and that must not pass silently.""" + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + for name in ("GOOGLE_CLIENT_ID", "MICROSOFT_CLIENT_ID", "SAML_IDP_METADATA_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("GENERIC_CLIENT_ID", "cid") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=True)) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "no usable id_token" in str(warnings[0]) + + +@pytest.mark.asyncio +async def test_no_warning_when_the_assertion_was_captured(monkeypatch, caplog): + from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import ( + assertion_from_sso_login, + ) + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + assertion = assertion_from_sso_login(_ema_id_token(), None) + assert assertion is not None + + retention_mock = AsyncMock(return_value=True) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(assertion, retention_enabled=retention_mock) + + assert _id_jag_gap_warnings(caplog) == [] + retention_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_no_warning_when_no_id_jag_server_is_registered(monkeypatch, caplog): + """Most deployments never register one; a warning about ID-JAG on every login there would + be pure noise and would train operators to ignore it.""" + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await warn_if_id_jag_assertion_uncaptured(None, retention_enabled=AsyncMock(return_value=False)) + + assert _id_jag_gap_warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_store_outage_does_not_break_the_login(monkeypatch, caplog): + from litellm.proxy.management_endpoints.ui_sso import ( + warn_if_id_jag_assertion_uncaptured, + ) + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + assert ( + await warn_if_id_jag_assertion_uncaptured( + None, retention_enabled=AsyncMock(side_effect=Exception("db down")) + ) + is None + ) + + assert _id_jag_gap_warnings(caplog) == [] + + +@pytest.mark.asyncio +async def test_browser_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog): + """Wiring: the browser login path must reach the diagnostic, not just define it.""" + monkeypatch.setenv("GOOGLE_CLIENT_ID", "cid") + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.cookies = {} + + with ( + patch( # test-quality-ok: endpoint test stubs the Prisma client lookup + "litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock() + ), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.user_custom_sso", None), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: endpoint reads proxy globals + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: endpoint reads proxy globals + patch( # test-quality-ok: endpoint test stubs key generation at its module boundary + "litellm.proxy.proxy_server.generate_key_helper_fn", + AsyncMock(return_value={"token": "sk-ui-key", "user_id": "canonical-user-id"}), + ), + patch( # test-quality-ok: endpoint test stubs the user database lookup + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=None), + ), + patch( # test-quality-ok: endpoint test stubs the admin database lookup + "litellm.proxy.management_endpoints.ui_sso.check_and_update_if_proxy_admin_id", + AsyncMock(return_value="internal_user"), + ), + patch( # test-quality-ok: endpoint test stubs assertion persistence + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + patch( # test-quality-ok: endpoint reads this module global without an injection seam + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=True), + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await SSOAuthenticationHandler.get_redirect_response_from_openid( + result=CustomOpenID( + id="raw-idp-subject", + email="u@example.com", + first_name="U", + last_name="Ser", + display_name="U Ser", + provider="google", + team_ids=[], + user_role=None, + ), + request=mock_request, + received_response=None, + generic_client_id=None, + ui_access_mode=None, + access_token_payload=None, + jwt_handler=None, + sso_assertion=None, + ) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "google" in str(warnings[0]) + + +@pytest.mark.asyncio +async def test_cli_funnel_reports_an_uncaptured_assertion(monkeypatch, caplog): + """Wiring: the CLI login path shares the gap, so it must share the diagnostic.""" + from litellm.proxy.management_endpoints.ui_sso import ( + _complete_cli_sso_callback_session, + ) + + monkeypatch.setenv("MICROSOFT_CLIENT_ID", "cid") + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + + user_info = MagicMock() + user_info.user_id = "cli-user-id" + user_info.user_role = "internal_user" + user_info.models = [] + user_info.teams = [] + + with ( + patch( # test-quality-ok: endpoint test stubs the user database lookup + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=user_info), + ), + patch( # test-quality-ok: endpoint test stubs CLI team lookup + "litellm.proxy.management_endpoints.ui_sso.fetch_cli_sso_team_details", + AsyncMock(return_value=[]), + ), + patch( # test-quality-ok: endpoint test stubs attribution metadata + "litellm.proxy.management_endpoints.ui_sso.build_cli_sso_attribution_metadata", + return_value={}, + ), + patch( # test-quality-ok: endpoint test stubs assertion persistence + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + patch( # test-quality-ok: endpoint reads this module global without an injection seam + "litellm.proxy.management_endpoints.ui_sso.ema_assertion_retention_enabled", + AsyncMock(return_value=True), + ), + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await _complete_cli_sso_callback_session( + request=mock_request, + key="cli-login-id", + flow={}, + result={"sub": "raw-idp-subject"}, + parsed_openid_result={ + "user_id": "raw-idp-subject", + "user_email": "u@example.com", + "user_role": None, + }, + user_defined_values=None, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + cli_sso_session_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + sso_assertion=None, + ) + + warnings = _id_jag_gap_warnings(caplog) + assert len(warnings) == 1 + assert "microsoft" in str(warnings[0]) + + def _cli_callback_kwargs(flow): return { "request": _cli_callback_request(),