From d36cdcb8f9808448eeedc99f50996feb1ea92536 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:30:21 +0000 Subject: [PATCH 01/10] fix(ui): show loading state instead of stale rows while a table search is pending Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../hooks/common/useResourceList.test.tsx | 22 +++++++++++++++++++ .../hooks/common/useResourceList.ts | 4 ++-- .../users/_components/view_users.test.tsx | 11 ++++++++++ .../users/_components/view_users.tsx | 2 +- .../components/TeamsPage/TeamsTable.test.tsx | 8 +++++++ .../src/components/TeamsPage/TeamsTable.tsx | 5 +++-- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 9 ++++++++ .../VirtualKeysPage/VirtualKeysTable.tsx | 5 +++-- 8 files changed, 59 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx index 3ca67b082ec..b0c478920f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx @@ -152,6 +152,28 @@ describe("useResourceList", () => { await waitFor(() => expect(lastCall().page_size).toBe(25)); }); + it("reports loading while a new search request is still pending", async () => { + let resolveSecond: ((value: ResourceListPage) => void) | undefined; + const fetchPage = vi.fn((query: ResourceListQuery) => { + calls.push(query); + if (calls.length === 1) return Promise.resolve(page([{ id: "a" }], 3)); + return new Promise>((resolve) => { + resolveSecond = resolve; + }); + }); + const { result } = renderList({ fetchPage }); + await waitFor(() => expect(result.current.rows).toEqual([{ id: "a" }])); + expect(result.current.isLoading).toBe(false); + + act(() => result.current.onSearchChange("zzz")); + await waitFor(() => expect(lastCall().q).toBe("zzz")); + expect(result.current.isLoading).toBe(true); + + act(() => resolveSecond?.(page([], 0))); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + expect(result.current.rows).toEqual([]); + }); + it("surfaces a failed page as an error instead of empty rows", async () => { const fetchPage = vi.fn(() => Promise.reject(new Error("boom"))); const { result } = renderList({ fetchPage }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts index 8a6376b2248..983e537bd00 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts @@ -86,7 +86,7 @@ export function useResourceList(options: UseResourceListOptions): Re enabled, placeholderData: (previous) => previous, }; - const { data, isLoading, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); + const { data, isLoading, isPlaceholderData, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); const toFirstPage = useCallback(() => setPagination((previous) => ({ ...previous, pageIndex: 0 })), []); @@ -123,7 +123,7 @@ export function useResourceList(options: UseResourceListOptions): Re return { rows, rowCount: data?.meta.total_count ?? 0, - isLoading, + isLoading: isLoading || isPlaceholderData, isFetching, error, refetch, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 6b0423067fd..10a13c1dbd7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -344,5 +344,16 @@ describe("ViewUserDashboard", () => { expect(latest[4]).toBeNull(); expect(latest[2]).toBe(1); }); + + it("replaces the previous rows with the loading state while the search request is pending", async () => { + renderDashboard(); + expect(await screen.findByText("test@example.com")).toBeInTheDocument(); + + userListCall.mockReturnValue(new Promise(() => undefined)); + fireEvent.change(screen.getByPlaceholderText("Search by email or ID…"), { target: { value: "zzznomatch" } }); + + expect(await screen.findByText("Loading users…")).toBeInTheDocument(); + expect(screen.queryByText("test@example.com")).not.toBeInTheDocument(); + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 1ed23e7f523..20ce22b6444 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -295,7 +295,7 @@ const ViewUserDashboard: React.FC = ({ { expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); }); +it("replaces the previous rows with the loading state while a new search is pending", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { isPlaceholderData: true, isFetching: true })); + renderTable(); + + expect(screen.getByText("Loading teams...")).toBeInTheDocument(); + expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); +}); + describe("sort contract – only backend-sortable columns are sortable", () => { it("requests the default created_at descending sort on first render", () => { renderTable(); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index 3fb19e522f3..a1ec61daa32 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -83,7 +83,8 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet const { data: teamsResponse, - isPending: isLoading, + isPending, + isPlaceholderData, isFetching, refetch, } = useTeamsTable(tablePagination.pageIndex + 1, tablePagination.pageSize, teamListOptions); @@ -161,7 +162,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet onColumnFiltersChange={handleColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" - isLoading={isLoading} + isLoading={isPending || isPlaceholderData} loadingMessage="Loading teams..." noDataMessage="No teams found" fillHeight diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 881b0b93ff9..bc9017468aa 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -289,6 +289,15 @@ it("should show a loading state on the initial load and hide the data", () => { expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); }); +it("replaces the previous rows with the loading state while a new search is pending", () => { + mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { isPlaceholderData: true, isFetching: true })); + + renderWithProviders(); + + expect(screen.getByText("Loading keys...")).toBeInTheDocument(); + expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); +}); + it("should show 'No keys found' message when the key list is empty", () => { mockUseKeys.mockReturnValue(keysResult([])); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index f28ae27b6bc..1394b7259c7 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -128,7 +128,8 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const { data: keys, - isPending: isLoading, + isPending, + isPlaceholderData, isFetching, refetch, } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); @@ -280,7 +281,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { onColumnFiltersChange={handleColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" - isLoading={isLoading} + isLoading={isPending || isPlaceholderData} loadingMessage="Loading keys..." noDataMessage="No keys found" fillHeight From f5f81e973aa2804b1c4119109973c9da48398f54 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:43:55 -0700 Subject: [PATCH 02/10] fix(budgets): preserve explicit reset interval clears --- .../budget_management_endpoints.py | 5 +++++ .../internal_user_endpoints.py | 10 ++++++---- .../tag_management_endpoints.py | 1 + litellm/proxy/management_helpers/utils.py | 11 ++++++++++- .../test_budget_endpoints.py | 10 +++++----- .../test_internal_user_endpoints.py | 16 ++++++++++++++++ .../components/AccessGroupBudgetModal.tsx | 6 +++++- .../_components/components/CreateTagModal.tsx | 6 +++++- .../tag-management/_components/tag_info.tsx | 2 +- ui/litellm-dashboard/src/components/Teams.tsx | 2 +- .../budget_duration_dropdown.tsx | 8 ++------ .../components/organisms/create_key_button.tsx | 2 +- .../src/components/tag_management/types.tsx | 2 +- .../src/components/team/EditMembership.tsx | 2 +- .../src/components/team/TeamInfo.tsx | 4 +++- 15 files changed, 63 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 62a24109dbb..81a607aaa43 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -14,6 +14,7 @@ All /budget management endpoints #### BUDGET TABLE MANAGEMENT #### import math from collections.abc import Mapping +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException @@ -176,6 +177,10 @@ async def update_budget( recomputed_reset_at: Final = ( {"budget_reset_at": get_budget_reset_time(budget_duration=budget_obj.budget_duration)} if budget_obj.budget_duration is not None and "budget_reset_at" not in budget_obj.model_fields_set + else MappingProxyType({"budget_reset_at": None}) + if "budget_duration" in budget_obj.model_fields_set + and budget_obj.budget_duration is None + and "budget_reset_at" not in budget_obj.model_fields_set else {} ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 10c11119006..000b7f874ee 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1254,8 +1254,8 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set() for k, v in data_json.items(): - if k == "max_budget": - if "max_budget" in fields_set: + if k in ("max_budget", "budget_duration"): + if k in fields_set: non_default_values[k] = v elif k == "model_max_budget": if k in fields_set: @@ -1283,8 +1283,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time validate_budget_duration(non_default_values["budget_duration"]) - non_default_values["budget_reset_at"] = get_budget_reset_time( - budget_duration=non_default_values["budget_duration"] + non_default_values["budget_reset_at"] = ( + get_budget_reset_time(budget_duration=non_default_values["budget_duration"]) + if non_default_values["budget_duration"] is not None + else None ) if "max_budget" not in non_default_values: diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index b74aa1a4e16..ab33d4bd766 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -438,6 +438,7 @@ async def update_tag( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, litellm_proxy_admin_name=litellm_proxy_admin_name, + budget_duration_cleared="budget_duration" in tag.model_fields_set and tag.budget_duration is None, ) # Get model names for model_info diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index e2d7262fb69..f3bd4b0f6dd 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -3,6 +3,7 @@ from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime from functools import wraps +from types import MappingProxyType from typing import Any, Final, Protocol from fastapi import HTTPException, Request @@ -180,6 +181,7 @@ async def handle_budget_for_entity( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, litellm_proxy_admin_name: str, + budget_duration_cleared: bool = False, ) -> str | None: """ Common helper to handle budget creation/updates for entities (organizations, tags, etc). @@ -208,7 +210,14 @@ async def handle_budget_for_entity( # Extract budget fields from data _json_data: Final = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data - _budget_data: Final = {k: v for k, v in _json_data.items() if k in budget_params} + _budget_data: Final = MappingProxyType( + { + k: _json_data.get(k) + for k in budget_params + if k in _json_data + or (k == "budget_duration" and existing_budget_id is not None and budget_duration_cleared) + } + ) # Check if budget_id is explicitly provided in the data data_budget_id: Final[str | None] = getattr(data, "budget_id", None) diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index 79d62f772bd..4b6815d7552 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -340,7 +340,8 @@ async def test_update_budget_recomputes_reset_at_when_duration_changes( @pytest.mark.asyncio -async def test_update_budget_preserves_explicit_reset_at(client_and_mocks): +@pytest.mark.parametrize("budget_duration", ["1d", None]) +async def test_update_budget_preserves_explicit_reset_at(client_and_mocks, budget_duration): """An explicit budget_reset_at from the caller always wins over recompute.""" client, _, mock_table = client_and_mocks captured = _capture_update_data(mock_table) @@ -350,7 +351,7 @@ async def test_update_budget_preserves_explicit_reset_at(client_and_mocks): "/budget/update", json={ "budget_id": "budget_explicit_reset", - "budget_duration": "1d", + "budget_duration": budget_duration, "budget_reset_at": explicit.isoformat(), }, ) @@ -377,8 +378,7 @@ async def test_update_budget_without_duration_leaves_reset_at_untouched( @pytest.mark.asyncio -async def test_update_budget_duration_none_does_not_recompute(client_and_mocks): - """Clearing budget_duration (explicit null) must not recompute against a None duration.""" +async def test_update_budget_duration_none_clears_obsolete_reset(client_and_mocks): client, _, mock_table = client_and_mocks captured = _capture_update_data(mock_table) @@ -389,7 +389,7 @@ async def test_update_budget_duration_none_does_not_recompute(client_and_mocks): assert resp.status_code == 200, resp.text assert "budget_duration" in captured and captured["budget_duration"] is None - assert "budget_reset_at" not in captured + assert captured["budget_reset_at"] is None @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 92b1ab1586d..0d8b19345f1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2097,6 +2097,22 @@ def test_update_internal_user_params_reset_max_budget_with_none(): assert non_default_values["user_id"] == "test_user" +def test_update_internal_user_params_explicit_duration_clear_overrides_role_default(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "internal_user_budget_duration", "30d") + data = UpdateUserRequest( + user_id="duration-clear-test", + user_role=LitellmUserRoles.INTERNAL_USER, + budget_duration=None, + ) + + updated = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) + + assert updated["budget_duration"] is None + assert updated["budget_reset_at"] is None + + def test_update_internal_user_params_ignores_other_nones(): """ Test that other fields are still filtered out if None diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx index 8df2a6620f9..cf5cf78d08a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx @@ -98,7 +98,11 @@ const AccessGroupBudgetModal: React.FC = ({ )} > {({ id, value, onChange }) => ( - + onChange(next ?? undefined)} + /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx index 6f6bdcb6fe4..adb5a167c35 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx @@ -143,7 +143,11 @@ const CreateTagModal: React.FC = ({ visible, onCancel, onSu )} > {({ id, value, onChange }) => ( - + onChange(next ?? undefined)} + /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx index e8cb358c0cf..1648a99bb0e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx @@ -28,7 +28,7 @@ const tagEditShape = { description: z.string().optional(), models: z.array(z.string()).optional(), max_budget: z.union([z.string(), z.number()]).optional(), - budget_duration: z.string().optional(), + budget_duration: z.string().nullish(), }; const tagEditSchema = z.object(tagEditShape); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index f0d21cc5350..dc531ea5dad 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -807,7 +807,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser showNeverResets placeholder={budgetDurationPlaceholder} value={value} - onChange={onChange} + onChange={(next) => onChange(next ?? undefined)} /> )} diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 3c38907c597..40ee857634c 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -14,7 +14,7 @@ const DURATION_LABELS: Record = { interface BudgetDurationDropdownProps { id?: string; value?: string | null; - onChange?: (value: string | undefined) => void; + onChange?: (value: string | null) => void; className?: string; style?: React.CSSProperties; placeholder?: string; @@ -31,11 +31,7 @@ const BudgetDurationDropdown: React.FC = ({ showNeverResets = false, }) => { return ( - diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 82580fd4667..b5789101f77 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1021,7 +1021,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp value={control.value as string | null | undefined} showNeverResets placeholder="Not set" - onChange={control.onChange} + onChange={(next) => control.onChange(next ?? undefined)} /> )} diff --git a/ui/litellm-dashboard/src/components/tag_management/types.tsx b/ui/litellm-dashboard/src/components/tag_management/types.tsx index 3cf17545fd6..88dfa28204d 100644 --- a/ui/litellm-dashboard/src/components/tag_management/types.tsx +++ b/ui/litellm-dashboard/src/components/tag_management/types.tsx @@ -41,7 +41,7 @@ export interface TagUpdateRequest { soft_budget?: number; tpm_limit?: number; rpm_limit?: number; - budget_duration?: string; + budget_duration?: string | null; } export interface TagDeleteRequest { diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index 036c4f1cc0b..909b5d56c97 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -157,7 +157,7 @@ const MemberModal = ({ onChange(next)} + onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)} /> ); default: diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 8e3a17c2622..ffc83d0165e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1467,7 +1467,9 @@ const TeamInfoView: React.FC = ({ showNeverResets placeholder="Inherit team reset period" value={value === null ? NEVER_RESETS_BUDGET_DURATION : value} - onChange={(next) => onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next)} + onChange={(next) => + onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next ?? undefined) + } /> )} From ad966d834020368e5ee1aa9974bb3ff1b5492cd1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:43:55 -0700 Subject: [PATCH 03/10] fix(projects): persist explicit budget cap clears --- .../management_endpoints/project_endpoints.py | 5 ++++- .../test_project_endpoints_prisma.py | 15 +++++++++++++++ .../hooks/projects/useUpdateProject.ts | 2 +- .../EditProjectModal.integration.test.tsx | 14 ++++++++++++++ .../ProjectModals/EditProjectModal.tsx | 2 +- .../_components/ProjectModals/ProjectBaseForm.tsx | 15 +++++++++++++-- .../ProjectModals/projectFormSchema.ts | 2 +- .../ProjectModals/projectFormUtils.test.ts | 6 +++--- .../_components/ProjectModals/projectFormUtils.ts | 12 ++++++++++-- 9 files changed, 62 insertions(+), 11 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index b2eda76f9ae..f40ced302ce 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -780,7 +780,10 @@ async def update_project( # Handle budget updates budget_fields = LiteLLM_BudgetTable.model_fields.keys() - budget_updates = {k: v for k, v in update_data.items() if k in budget_fields} + budget_updates = { + **{k: v for k, v in update_data.items() if k in budget_fields}, + **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}), + } if budget_updates and existing_project.budget_id: # Update existing budget diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c23b203feba..36878fa698c 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1292,6 +1292,21 @@ async def test_update_project_leaves_metadata_untouched_when_no_limit_is_sent(mo assert "metadata" not in _written_project_data(mock_prisma) +@pytest.mark.asyncio +async def test_update_project_clears_only_the_explicit_budget_cap(monkeypatch): + mock_prisma = _project_update_mocks(monkeypatch, {}) + mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-clear-test" + mock_prisma.db.litellm_budgettable.update = mock.AsyncMock() + + await _run_project_update("project-clear-test", max_budget=None) + + mock_prisma.db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-clear-test"}, + data={"max_budget": None, "updated_by": "1234"}, + ) + assert "max_budget" not in _written_project_data(mock_prisma) + + @pytest.mark.parametrize("entry", ["all-proxy-models", "*", "azure/*"]) def test_enforce_project_model_quota_rejects_entries_that_expand_at_request_time(entry): """A quota keyed on a wildcard entry is never applied by the limiter, so it fails loudly.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts index 8e6bad04a28..e6dec85128c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts @@ -10,7 +10,7 @@ export interface ProjectUpdateParams { description?: string; team_id?: string; models?: string[]; - max_budget?: number; + max_budget?: number | null; blocked?: boolean; guardrails?: string[]; metadata?: Record; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx index 5b84aa15dd1..2e39d45e660 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx @@ -102,6 +102,20 @@ describe("EditProjectModal submit payload", () => { }); }); + it("should send an explicit clear after blanking a saved budget", async () => { + const user = setup(); + renderModal(); + + const budgetInput = screen.getByRole("spinbutton", { name: "Max Budget (USD)" }); + await user.clear(budgetInput); + await user.tab(); + expect(budgetInput).toHaveValue(null); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(JSON.parse(JSON.stringify(variables().params))).toMatchObject({ max_budget: null }); + }); + it("includes the advanced fields once Advanced Settings has been opened, even after collapsing it again", async () => { const user = setup(); renderModal(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 77f28b05ea5..31582da5b14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -86,7 +86,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit onChange(toOptionalNumber(event.target.value))} + value={Number.isNaN(value) ? "" : value ?? ""} + onInput={(event) => { + if (event.currentTarget.validity.badInput || Number.isNaN(value)) { + onChange( + event.currentTarget.validity.badInput + ? Number.NaN + : toOptionalNumber(event.currentTarget.value) ?? null, + ); + } + }} + onChange={(event) => + onChange(event.target.validity.badInput ? Number.NaN : toOptionalNumber(event.target.value) ?? null) + } /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts index d4c85d89616..6e6318a9609 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts @@ -22,7 +22,7 @@ export const projectFormSchema = z .pipe(z.string({ error: "Please select a team" }).min(1, "Please select a team")), description: z.string().optional(), models: z.array(z.string()), - max_budget: z.number().optional(), + max_budget: z.number().nullish(), isBlocked: z.boolean(), guardrails: z.array(z.string()).optional(), modelLimits: z.array(modelLimitSchema).optional(), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts index 28d791ddf54..e76555b3ca3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts @@ -27,9 +27,9 @@ describe("buildProjectCreateParams", () => { expect(result.description).toBe("A description"); }); - it("should pass through max_budget when provided", () => { - const result = buildProjectCreateParams({ ...baseValues, max_budget: 50.0 }); - expect(result.max_budget).toBe(50.0); + it.each([50.0, 1e308])("should preserve a finite max_budget of %s", (maxBudget) => { + const result = buildProjectCreateParams({ ...baseValues, max_budget: maxBudget }); + expect(JSON.parse(JSON.stringify(result)).max_budget).toBe(maxBudget); }); it("should build model_rpm_limit from modelLimits entries", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts index 71c1f5c79af..60e97939e69 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts @@ -16,6 +16,11 @@ const buildModelLimitMap = ( const buildMetadata = (entries: ProjectFormValues["metadata"]): Record | undefined => entries && Object.fromEntries(entries.flatMap((entry) => (entry.key ? [[entry.key, entry.value] as const] : []))); +const roundBudget = (value: number): number => { + const rounded = Math.round(value * 100) / 100; + return Number.isFinite(rounded) ? rounded : value; +}; + const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => { const limitEntries = values.modelLimits ?? []; const modelRpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.rpm); @@ -35,7 +40,7 @@ const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => project_alias: values.project_alias, description: values.description, models: values.models ?? [], - max_budget: values.max_budget === undefined ? undefined : Math.round(values.max_budget * 100) / 100, + max_budget: values.max_budget == null ? undefined : roundBudget(values.max_budget), blocked: values.isBlocked ?? false, ...guardrailsParam, ...(keep(modelRpmLimit) && { model_rpm_limit: modelRpmLimit }), @@ -53,4 +58,7 @@ export const buildProjectCreateParams = (values: ProjectFormValues) => buildProj * /project/update leaves an omitted key untouched, so a limit the operator cleared has to go out as * an explicitly empty map. Omitting it is what silently kept a removed quota enforced. */ -export const buildProjectUpdateParams = (values: ProjectFormValues) => buildProjectApiParams(values, true); +export const buildProjectUpdateParams = (values: ProjectFormValues, savedMaxBudget?: number | null) => ({ + ...buildProjectApiParams(values, true), + ...(values.max_budget == null && savedMaxBudget != null ? { max_budget: null } : {}), +}); From b24480941bd06cb615f0489767d35f1616a775d0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:43:55 -0700 Subject: [PATCH 04/10] fix(ui): retain cleared user budget values after save --- .../user_info_view.integration.test.tsx | 31 ++++++++++++++++++- .../_components/view_users/user_info_view.tsx | 5 +-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx index 69254b1ffe4..0f1a44851c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; import UserInfoView from "./user_info_view"; @@ -163,6 +163,35 @@ describe("UserInfoView add-to-team form", () => { expect(await openEditor(user)).toHaveValue(42); }); + + it("should keep Unlimited selected after saving and reopening the user", async () => { + const user = setup(); + render(); + + await openEditor(user); + await user.click(screen.getByRole("checkbox", { name: "Unlimited Budget" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); + expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ max_budget: null }); + await openEditor(user); + expect(screen.getByRole("checkbox", { name: "Unlimited Budget" })).toBeChecked(); + }); + + it("should keep a cleared reset period after saving and reopening the user", async () => { + const user = setup(); + render(); + + await openEditor(user); + await user.click(screen.getByRole("combobox", { name: "Reset Budget" })); + await user.click(await screen.findByRole("option", { name: "n/a" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); + expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ budget_duration: null }); + await openEditor(user); + expect(screen.getByRole("combobox", { name: "Reset Budget" })).toHaveTextContent("n/a"); + }); }); it("offers only the teams the user is not already a member of", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index eed39c8e585..e083e549552 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -332,8 +332,9 @@ export default function UserInfoView({ user_email: formValues.user_email ?? userData.user_email, user_alias: formValues.user_alias ?? userData.user_alias, models: formValues.models ?? userData.models, - max_budget: formValues.max_budget ?? userData.max_budget, - budget_duration: formValues.budget_duration ?? userData.budget_duration, + max_budget: formValues.max_budget === undefined ? userData.max_budget : formValues.max_budget, + budget_duration: + formValues.budget_duration === undefined ? userData.budget_duration : formValues.budget_duration, metadata: formValues.metadata ?? userData.metadata, model_max_budget: formValues.model_max_budget ?? userData.model_max_budget, object_permission: mcpEntitlement From 261807114a50e09b7fe6b48a7d8709dbf25e9c02 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:10:58 -0700 Subject: [PATCH 05/10] fix(proxy): accept both deferred stream logging arg shapes on native routes (#40869) * fix(proxy): accept both deferred stream logging arg shapes on native routes _arm_deferred_stream_dispatch armed a one-argument closure on every anthropic_messages/aresponses stream that was not a CustomStreamWrapper or a LiteLLMCompletionStreamingIterator. The bridged /v1/messages path returns a plain SSE generator that shares its inner CustomStreamWrapper logging_obj, so it stores (assembled_response, cache_hit) and _fire_deferred_stream_logging raised TypeError, dropping spend logs and callbacks and ending the stream with an error. The closure now dispatches on the stored args shape: a single coroutine is enqueued, a two-tuple runs success handlers, anything else is logged and dropped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert dropped deferred payload via caplog instead of patching the logger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 66 ++++++++-------- .../test_deferred_guardrail_logging.py | 78 ++++++++++++++++++- 2 files changed, 107 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 289fe086379..e6ed60ba177 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,7 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -3210,20 +3210,24 @@ class ProxyBaseLLMRequestProcessing: end-of-stream blocks complete, so the spend log sees guardrail_information. - Three closure shapes, matching who owns logging for the stream: + Two closure shapes, matching who owns logging for the stream: - CustomStreamWrapper (chat completions) stores (assembled_response, cache_hit); the closure also runs non-apply_guardrail post-call hooks via _run_deferred_stream_guardrails. - - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares - its inner CustomStreamWrapper's logging_obj, so it stores the same - (assembled_response, cache_hit) shape; the closure only dispatches - success logging, matching the route's pre-existing hook surface. - - Native anthropic_messages/aresponses iterators store a single - ready-made logging coroutine to enqueue. + - Every other anthropic_messages/aresponses stream gets a closure + that dispatches on the stored args shape, because the arming site + cannot tell the producers apart: native iterators store a single + ready-made logging coroutine to enqueue, while bridged streams + (LiteLLMCompletionStreamingIterator, and the plain SSE generator + AnthropicStreamWrapper returns for bridged /v1/messages) share + their inner CustomStreamWrapper's logging_obj and so store + (assembled_response, cache_hit); for those the closure only + dispatches success logging, matching the route's pre-existing + hook surface. - Raw async generators from passthrough routes bypass all three and - would orphan the closure, so they are not armed here. + Raw async generators from passthrough routes bypass both and would + orphan the closure, so they are not armed here. The router wraps iterators that cannot carry _hidden_params in HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the @@ -3257,31 +3261,27 @@ class ProxyBaseLLMRequestProcessing: if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): return - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - - if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): - _captured_bridge_logging_obj: Final = logging_obj - - async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: - await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete - return - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + _captured_native_logging_obj: Final = logging_obj + + async def _on_deferred_native_stream_complete(*args: object) -> None: + match args: + case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + case (assembled_response, cache_hit): + await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + case _: + verbose_proxy_logger.error( + "Deferred stream logging dropped: unexpected stored args shape %s", + tuple(type(arg).__name__ for arg in args), + ) logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 11c3d2f8b20..c550a0a41d2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -15,6 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. """ import asyncio +import logging from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -1422,7 +1423,7 @@ class TestArmDeferredStreamDispatch: async def test_native_stream_closure_enqueues_single_coroutine(self): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, _ = self._dispatch_recording_logging_obj() + logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): yield b"x" @@ -1433,20 +1434,89 @@ class TestArmDeferredStreamDispatch: user_api_key_dict=MagicMock(), logging_obj=logging_obj, ) - closure = logging_obj._on_deferred_stream_complete - assert closure is not None + assert logging_obj._on_deferred_stream_complete is not None async def _logging_coroutine(): return None coro = _logging_coroutine() + logging_obj._deferred_stream_complete_args = (coro,) with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" ) as mock_enqueue: - await closure(coro) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) mock_enqueue.assert_called_once_with(async_coroutine=coro) + assert recorded == {} coro.close() + @pytest.mark.asyncio + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_raw_generator_stream_storing_csw_arg_shape_dispatches_success(self, route_type): + """Bridged /v1/messages returns AnthropicStreamWrapper's plain SSE + generator, which shares its inner CustomStreamWrapper's logging_obj and + so stores (assembled_response, cache_hit). The closure armed for a raw + generator must accept that shape too, or _fire_deferred_stream_logging + raises TypeError and the request loses its spend log and callbacks.""" + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type=route_type, + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, True) + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded["result"] is assembled + assert recorded["cache_hit"] is True + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())]) + async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + logging_obj._deferred_stream_complete_args = stored_args + with ( + patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue, + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + ): + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded == {} + dropped = [r for r in caplog.records if r.getMessage().startswith("Deferred stream logging dropped")] + assert len(dropped) == 1 + @pytest.mark.asyncio async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper From 09d63b259c515fbb2f772e5aea5d259eb3c6ef41 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 12 Sep 2026 21:11:02 +0000 Subject: [PATCH 06/10] test(fireworks): drop hardcoded price snapshot tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 248 ------------------------------- 1 file changed, 248 deletions(-) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 19ed31c7b22..965f960f49b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4245,254 +4245,6 @@ def test_deepseek_flash_completion_cost(): assert cost == pytest.approx(1.50, abs=1e-9) -_FIREWORKS_MODELS = [ - ( - "accounts/fireworks/models/glm-5p2", - 1.4e-06, - 4.4e-06, - 1.4e-07, - 1048576, - 131072, - False, - True, - ), - ( - "accounts/fireworks/models/glm-5p1", - 1.4e-06, - 4.4e-06, - 2.6e-07, - 202800, - 131072, - False, - True, - ), - ( - "accounts/fireworks/routers/glm-5p1-fast", - 2.8e-06, - 8.8e-06, - 5.2e-07, - 202800, - 131072, - False, - True, - ), - ( - "accounts/fireworks/models/qwen3p7-plus", - 4e-07, - 1.6e-06, - 8e-08, - 262144, - 65536, - True, - True, - ), - ( - "accounts/fireworks/models/minimax-m3", - 3e-07, - 1.2e-06, - 6e-08, - 512000, - 512000, - True, - True, - ), - ( - "accounts/fireworks/models/minimax-m2p7", - 3e-07, - 1.2e-06, - 6e-08, - 196608, - 196608, - False, - True, - ), - ( - "accounts/fireworks/models/kimi-k2p7-code", - 9.5e-07, - 4e-06, - 1.9e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/routers/kimi-k2p7-code-fast", - 1.9e-06, - 8e-06, - 3.8e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/models/kimi-k2p6", - 9.5e-07, - 4e-06, - 1.6e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/routers/kimi-k2p6-fast", - 2e-06, - 8e-06, - 3e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/models/gpt-oss-120b", - 1.5e-07, - 6e-07, - 1.5e-08, - 131072, - 32768, - False, - True, - ), - ( - "accounts/fireworks/models/gpt-oss-20b", - 7e-08, - 3e-07, - 3.5e-08, - 131072, - 32768, - False, - True, - ), - ( - "accounts/fireworks/models/deepseek-v4-pro", - 1.74e-06, - 3.48e-06, - 1.45e-07, - 1048576, - 384000, - False, - True, - ), - ( - "accounts/fireworks/models/deepseek-v4-flash", - 1.4e-07, - 2.8e-07, - 2.8e-08, - 1048576, - 384000, - False, - True, - ), -] - -_FIREWORKS_SHORT_FORMS = [ - "glm-5p2", - "glm-5p1", - "qwen3p7-plus", - "minimax-m3", - "minimax-m2p7", - "kimi-k2p7-code", - "kimi-k2p6", - "gpt-oss-120b", - "gpt-oss-20b", - "deepseek-v4-pro", - "deepseek-v4-flash", -] - -_FIREWORKS_ROUTER_SHORT_FORMS = [ - "glm-5p1-fast", - "kimi-k2p6-fast", - "kimi-k2p7-code-fast", -] - - -def _assert_fireworks_entry( - model_cost, - model_path, - expected_input, - expected_output, - expected_cache, - expected_max_input, - expected_max_output, - expected_vision, - expected_reasoning, -): - info = model_cost.get(f"fireworks_ai/{model_path}") - assert info is not None, f"fireworks_ai/{model_path} missing from model cost map" - assert info["litellm_provider"] == "fireworks_ai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache - assert info["max_input_tokens"] == expected_max_input - assert info["max_output_tokens"] == expected_max_output - assert info["max_tokens"] == expected_max_output - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info["supports_reasoning"] is expected_reasoning - assert info["supports_response_schema"] is True - assert info["supports_vision"] is expected_vision - - -def test_fireworks_models_in_cost_map(): - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - for entry in _FIREWORKS_MODELS: - _assert_fireworks_entry(model_cost, *entry) - - for short in _FIREWORKS_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/models/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - for short in _FIREWORKS_ROUTER_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - -def test_fireworks_models_in_backup_cost_map(): - import json - from pathlib import Path - - json_path = ( - Path(__file__).parents[2] - / "litellm" - / "model_prices_and_context_window_backup.json" - ) - with open(json_path) as f: - model_cost = json.load(f) - - for entry in _FIREWORKS_MODELS: - _assert_fireworks_entry(model_cost, *entry) - - for short in _FIREWORKS_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/models/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - for short in _FIREWORKS_ROUTER_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - class TestBedrockBaseModelLabelKeepsTools: """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly label must not silently drop ``tools``/``tool_choice`` under ``drop_params``.""" From 311d9bba37b44099461e6e699232b0903474fc62 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:19:20 -0700 Subject: [PATCH 07/10] perf(policy_engine): dedup attachments in one pass after sorting (#40883) get_attached_policies_with_reasons rescanned the sorted matches with next() once per distinct policy, which is quadratic and misses the one second budget past a few thousand global attachments. Build a policy to broadest attachment map in one pass instead, keeping the specificity sort and result order. Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/attachment_registry.py | 6 +++++- .../policy_engine/test_attachment_registry.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index a8ead86ac36..76b2291774e 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -6,6 +6,7 @@ This allows the same policy to be attached to multiple scopes. """ from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict from litellm._logging import verbose_proxy_logger @@ -141,8 +142,11 @@ class AttachmentRegistry: ), key=_attachment_specificity, ) + broadest_attachment_by_policy: Final = MappingProxyType( + {attachment.policy: attachment for attachment in reversed(matching_attachments)} + ) unique_attachments: Final = tuple( - next(attachment for attachment in matching_attachments if attachment.policy == policy_name) + broadest_attachment_by_policy[policy_name] for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 87b1bc56659..fa37a02a37c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -4,6 +4,7 @@ Unit tests for AttachmentRegistry - tests policy attachment matching. Tests the main entry point: get_attached_policies() """ +import time from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock @@ -222,6 +223,21 @@ class TestGetAttachedPolicies: # Should only appear once assert attached.count("multi-policy") == 1 + def test_many_distinct_policies_resolve_in_linear_time(self): + policy_count = 20_000 + registry = AttachmentRegistry() + registry.load_attachments( + [{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)] + ) + context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") + + started = time.perf_counter() + attached = registry.get_attached_policies(context) + elapsed = time.perf_counter() - started + + assert attached == [f"policy-{index}" for index in range(policy_count)] + assert elapsed < 1.0, f"{policy_count} attachments took {elapsed:.2f}s, dedup is no longer one pass" + def test_no_attachments_returns_empty(self): """Test empty attachments returns empty list.""" registry = AttachmentRegistry() From fbcc6021220e54460bd27792dd1426e189abae61 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 12 Sep 2026 21:19:35 +0000 Subject: [PATCH 08/10] test(fireworks): stop pinning prices in the cost-map tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 203 +++++++++++++++++++++++++++++++ 1 file changed, 203 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 965f960f49b..835e87aff88 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4245,6 +4245,209 @@ def test_deepseek_flash_completion_cost(): assert cost == pytest.approx(1.50, abs=1e-9) +_FIREWORKS_MODELS = [ + ( + "accounts/fireworks/models/glm-5p2", + 1048576, + 131072, + False, + True, + ), + ( + "accounts/fireworks/models/glm-5p1", + 202800, + 131072, + False, + True, + ), + ( + "accounts/fireworks/routers/glm-5p1-fast", + 202800, + 131072, + False, + True, + ), + ( + "accounts/fireworks/models/qwen3p7-plus", + 262144, + 65536, + True, + True, + ), + ( + "accounts/fireworks/models/minimax-m3", + 512000, + 512000, + True, + True, + ), + ( + "accounts/fireworks/models/minimax-m2p7", + 196608, + 196608, + False, + True, + ), + ( + "accounts/fireworks/models/kimi-k2p7-code", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/routers/kimi-k2p7-code-fast", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/models/kimi-k2p6", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/routers/kimi-k2p6-fast", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/models/gpt-oss-120b", + 131072, + 32768, + False, + True, + ), + ( + "accounts/fireworks/models/gpt-oss-20b", + 131072, + 32768, + False, + True, + ), + ( + "accounts/fireworks/models/deepseek-v4-pro", + 1048576, + 384000, + False, + True, + ), + ( + "accounts/fireworks/models/deepseek-v4-flash", + 1048576, + 384000, + False, + True, + ), +] + +_FIREWORKS_SHORT_FORMS = [ + "glm-5p2", + "glm-5p1", + "qwen3p7-plus", + "minimax-m3", + "minimax-m2p7", + "kimi-k2p7-code", + "kimi-k2p6", + "gpt-oss-120b", + "gpt-oss-20b", + "deepseek-v4-pro", + "deepseek-v4-flash", +] + +_FIREWORKS_ROUTER_SHORT_FORMS = [ + "glm-5p1-fast", + "kimi-k2p6-fast", + "kimi-k2p7-code-fast", +] + + +def _assert_fireworks_entry( + model_cost, + model_path, + expected_max_input, + expected_max_output, + expected_vision, + expected_reasoning, +): + info = model_cost.get(f"fireworks_ai/{model_path}") + assert info is not None, f"fireworks_ai/{model_path} missing from model cost map" + assert info["litellm_provider"] == "fireworks_ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert "cache_read_input_token_cost" in info + assert info["max_input_tokens"] == expected_max_input + assert info["max_output_tokens"] == expected_max_output + assert info["max_tokens"] == expected_max_output + assert info["supports_function_calling"] is True + assert info["supports_tool_choice"] is True + assert info["supports_reasoning"] is expected_reasoning + assert info["supports_response_schema"] is True + assert info["supports_vision"] is expected_vision + + +def test_fireworks_models_in_cost_map(): + import json + from pathlib import Path + + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + model_cost = json.load(f) + + for entry in _FIREWORKS_MODELS: + _assert_fireworks_entry(model_cost, *entry) + + for short in _FIREWORKS_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/models/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + for short in _FIREWORKS_ROUTER_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + +def test_fireworks_models_in_backup_cost_map(): + import json + from pathlib import Path + + json_path = ( + Path(__file__).parents[2] + / "litellm" + / "model_prices_and_context_window_backup.json" + ) + with open(json_path) as f: + model_cost = json.load(f) + + for entry in _FIREWORKS_MODELS: + _assert_fireworks_entry(model_cost, *entry) + + for short in _FIREWORKS_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/models/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + for short in _FIREWORKS_ROUTER_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + class TestBedrockBaseModelLabelKeepsTools: """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly label must not silently drop ``tools``/``tool_choice`` under ``drop_params``.""" From 5b36de4646451da7cec224ab481d676315d4ffcb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:39:40 -0700 Subject: [PATCH 09/10] fix(guardrails): log mask when a guardrail adds request keys (#40882) * fix(guardrails): log mask when a guardrail adds request keys _inputs_were_modified only compared keys present in the pre-hook baseline, so a guardrail that injected a new key such as tools was logged as allow. Compare over the union of both key sets, and narrow the pre_call return value to the same prompt-bearing keys the baseline holds so passthrough stays allow. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): snapshot apply_guardrail inputs before the hook mutates them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 29 +++++---- .../guardrail_hooks/azure/prompt_shield.py | 1 + .../integrations/test_custom_guardrail.py | 59 +++++++++++++++++++ 3 files changed, 79 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8a976a966a6..39adea30828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -5,6 +5,7 @@ import os import secrets from collections.abc import Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args from litellm._logging import verbose_logger @@ -1279,6 +1280,7 @@ class CustomGuardrail(CustomLogger): guardrail_response: Final = self._summarize_guardrail_response( response=response, original_inputs=original_inputs, + event_type=event_type, ) verbose_logger.debug("Guardrail response: %s", response) @@ -1298,6 +1300,7 @@ class CustomGuardrail(CustomLogger): self, response: object, original_inputs: Mapping[str, object] | None, + event_type: GuardrailEventHooks | None, ) -> object: """Reduce a hook's return value to what is safe to log as ``guardrail_response``. @@ -1305,15 +1308,21 @@ class CustomGuardrail(CustomLogger): returns the (possibly modified) request payload. Neither is a provider verdict, and logging them verbatim ships the user's prompt to every logging sink (OTEL spans, Datadog, spend logs), so both collapse to ``"allow"`` / ``"mask"`` by comparing - against ``original_inputs``, a copy taken before the hook ran. A string result is the - hook's own rejection message (the proxy turns it into a 400), not user input, so it is - logged as is. + against ``original_inputs``, a copy taken before the hook ran. A pre_call baseline only + holds the prompt-bearing keys, so the returned request is narrowed to those same keys + before the comparison. A string result is the hook's own rejection message (the proxy + turns it into a 400), not user input, so it is logged as is. """ if response is None: return {} if original_inputs is None or not isinstance(response, Mapping): return response - return "mask" if self._inputs_were_modified(original_inputs, response) else "allow" + compared_response: Final[Mapping[str, object]] = ( + MappingProxyType({key: value for key, value in response.items() if key in _PRE_CALL_CONTENT_KEYS}) + if event_type == GuardrailEventHooks.pre_call + else response + ) + return "mask" if self._inputs_were_modified(original_inputs, compared_response) else "allow" @staticmethod def _is_guardrail_intervention(e: Exception) -> bool: @@ -1355,8 +1364,8 @@ class CustomGuardrail(CustomLogger): raise e def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool: - """True when any baseline key's value differs in ``response`` (mask), False otherwise (allow).""" - return any(response.get(key) != value for key, value in original_inputs.items()) + """True when any key of either mapping differs between them (mask), False otherwise (allow).""" + return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys()) def mask_content_in_string( self, @@ -1476,13 +1485,13 @@ def _original_inputs_for( ) -> dict | None: # mutable-ok: matches _process_response(original_inputs=) signature """Baseline the hook's return value is compared against to decide "allow" vs "mask". - ``apply_guardrail`` masks a fresh ``inputs`` dict, so that dict is the baseline. Pre-call - hooks edit the request in place and return it, so the baseline is a deep copy of the - prompt-bearing keys taken before the hook runs. + Hooks may edit their argument in place and return it, so the baseline is always a deep + copy taken before the hook runs: the whole ``inputs`` dict for ``apply_guardrail``, the + prompt-bearing request keys for pre-call hooks. """ if func_name == "apply_guardrail": inputs: Final = kwargs.get("inputs") - return inputs if isinstance(inputs, dict) else None + return copy.deepcopy(inputs) if isinstance(inputs, dict) else None if event_type != GuardrailEventHooks.pre_call: return None return {key: copy.deepcopy(value) for key, value in request_data.items() if key in _PRE_CALL_CONTENT_KEYS} diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index de9618a44a1..a0724b75ec7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -341,6 +341,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai guardrail_response: Final = self._summarize_guardrail_response( response=response, original_inputs=original_inputs, + event_type=event_type, ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index bb4822eae57..2fd5fa76d8e 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -3016,3 +3016,62 @@ class TestPreCallHookResponseIsNotLoggedVerbatim: ) assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_pre_call_hook_adding_tools_logs_mask(self): + class ToolInjectingGuardrail(CustomGuardrail): + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: object, + data: dict[str, object], + call_type: str, + ) -> dict[str, object]: + return {**data, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]} + + data = self._request() + await ToolInjectingGuardrail(guardrail_name="g").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion" + ) + + assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_apply_guardrail_adding_tools_logs_mask(self): + class ToolInjectingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]} + + data = self._request() + await ToolInjectingGuardrail(guardrail_name="g").apply_guardrail( + inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request" + ) + + assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_apply_guardrail_masking_inputs_in_place_logs_mask(self): + class InPlaceMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + inputs["texts"] = [""] + return inputs + + data = self._request() + await InPlaceMaskingGuardrail(guardrail_name="g").apply_guardrail( + inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request" + ) + + assert self._logged_response(data) == "mask" From 147eb23aaba9a2588ed5fee236d1953870ea74ef Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 14:40:01 -0700 Subject: [PATCH 10/10] bump: litellm-enterprise 0.1.66 -> 0.1.67, litellm-proxy-extras 0.4.96 -> 0.4.97 --- 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 903c5155a12..c049bf68c46 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.66" +version = "0.1.67" 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.66" +version = "0.1.67" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 7d4c78088f1..f94591872a4 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.96" +version = "0.4.97" 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.96" +version = "0.4.97" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index 31c498ccbbb..62ce4b4fd61 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.96", - "litellm-enterprise==0.1.66", + "litellm-proxy-extras==0.4.97", + "litellm-enterprise==0.1.67", "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 9c659be658f..eb4cdef76f1 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-08T03:56:24.358378Z" +exclude-newer = "2026-09-09T21:39:49.468411Z" exclude-newer-span = "P3D" [manifest] @@ -4773,12 +4773,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.66" +version = "0.1.67" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.96" +version = "0.4.97" source = { editable = "litellm-proxy-extras" } [[package]]