From 6a540a1bf848129dc16228d1b23f92120d7a7f03 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sat, 8 Aug 2026 20:17:12 -0700
Subject: [PATCH 1/6] fix(ui): gate organization and agent usage views behind
 capabilities

The Usage page admits internal users because their own usage view works,
but the entity breakdown selector inside it also offered Organization
Usage, so picking it fired /organization/daily/activity and collected a
401. Neither that route nor /agent/daily/activity appears in any non-admin
route list, so both are default-deny. The team breakdown leaked the second
one too: it fetches agent activity unconditionally to fill its Top Agents
card, which 401s for the same roles.

Adds viewOrganizationUsage and viewAgentUsage to the existing capability
map and points the selector option, the page section, and the fetch's
enabled flag at the same capability, so a role that cannot call the
endpoint never sees the breakdown and never issues the request. The team
and tag breakdowns, which internal users can read, are untouched, and the
default Usage view was already one of those.
---
 .../EntityUsage/EntityUsage.test.tsx          | 43 ++++++++++++++++++-
 .../components/EntityUsage/EntityUsage.tsx    | 34 +++++++++++----
 .../components/UsagePageView.test.tsx         | 25 +++++++++++
 .../_components/components/UsagePageView.tsx  |  9 ++--
 .../UsageViewSelect/UsageViewSelect.test.tsx  | 29 +++++++++++--
 .../UsageViewSelect/UsageViewSelect.tsx       | 20 +++++----
 .../src/utils/capabilities.test.ts            | 36 ++++++++++------
 .../src/utils/capabilities.ts                 |  2 +
 8 files changed, 160 insertions(+), 38 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
index 82ca66b10c0..11528117f1e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
@@ -1,4 +1,4 @@
-import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react";
+import { act, cleanup, fireEvent, render, screen, waitFor, within } from "@testing-library/react";
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import * as networking from "@/components/networking";
 import EntityUsage from "./EntityUsage";
@@ -856,6 +856,47 @@ describe("EntityUsage", () => {
     expect(logo.getAttribute("src")).toContain("openai_small");
   });
 
+  describe("capability gating", () => {
+    it.each([
+      ["organization", () => mockOrganizationDailyActivityCall, "Organization Spend Overview"],
+      ["agent", () => mockAgentDailyActivityCall, "Agent Spend Overview"],
+    ] as const)("fetches %s activity for an admin but not for an internal user", async (entityType, call, heading) => {
+      render(<EntityUsage {...defaultProps} entityType={entityType} />);
+      await waitFor(() => {
+        expect(call()).toHaveBeenCalled();
+      });
+
+      cleanup();
+      call().mockClear();
+
+      render(<EntityUsage {...defaultProps} entityType={entityType} userRole="Internal User" />);
+      expect(await screen.findByText(heading)).toBeInTheDocument();
+      expect(call()).not.toHaveBeenCalled();
+    });
+
+    it("keeps the team breakdown but drops its agent sub-fetch for an internal user", async () => {
+      render(<EntityUsage {...defaultProps} entityType="team" userRole="Internal User" />);
+
+      await waitFor(() => {
+        expect(mockTeamDailyActivityCall).toHaveBeenCalled();
+      });
+      expect(screen.getByText("Team Spend Overview")).toBeInTheDocument();
+
+      expect(mockAgentDailyActivityCall).not.toHaveBeenCalled();
+      expect(screen.queryByText("Agent Activity")).not.toBeInTheDocument();
+      expect(screen.queryByText("Top Agents Driving Spend")).not.toBeInTheDocument();
+    });
+
+    it("keeps the tag breakdown for an internal user", async () => {
+      render(<EntityUsage {...defaultProps} entityType="tag" userRole="Internal User" />);
+
+      await waitFor(() => {
+        expect(mockTagDailyActivityCall).toHaveBeenCalled();
+      });
+      expect(screen.getByText("Tag Spend Overview")).toBeInTheDocument();
+    });
+  });
+
   it("renders a letter avatar instead of an img for an unknown provider slug", async () => {
     const spendDataUnknownProvider = {
       ...mockSpendData,
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
index 4d44791d1a9..5a0f2abf15b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
@@ -2,6 +2,7 @@ import useTeams from "@/app/(dashboard)/hooks/useTeams";
 import { BarChart, DonutChart } from "@/components/shared/charts";
 import { MoneyCell } from "@/components/shared/table_cells";
 import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { hasCapability, type Capability } from "@/utils/capabilities";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import {
   Card,
@@ -108,7 +109,19 @@ const ENTITY_FETCH_FNS: Record<EntityType, (...args: any[]) => Promise<any>> = {
   user: userDailyActivityCall,
 };
 
-const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, entityId, entityList, dateValue }) => {
+const ENTITY_CAPABILITIES: Partial<Record<EntityType, Capability>> = {
+  organization: "viewOrganizationUsage",
+  agent: "viewAgentUsage",
+};
+
+const EntityUsage: React.FC<EntityUsageProps> = ({
+  accessToken,
+  entityType,
+  entityId,
+  entityList,
+  userRole,
+  dateValue,
+}) => {
   const { teams } = useTeams();
   const [selectedTags, setSelectedTags] = useState<string[]>([]);
   const [modelViewType, setModelViewType] = useState<ModelViewType>("groups");
@@ -125,7 +138,11 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
   }, [entityType, selectedTags]);
 
   const fetchFn = ENTITY_FETCH_FNS[entityType];
-  const enabled = !!accessToken && !!startTime && !!endTime;
+  const entityCapability = ENTITY_CAPABILITIES[entityType];
+  const canViewEntity = entityCapability === undefined || hasCapability(userRole, entityCapability);
+  const showAgentBreakdown = entityType === "team" && hasCapability(userRole, "viewAgentUsage");
+  const hasRequestWindow = !!accessToken && !!startTime && !!endTime;
+  const enabled = hasRequestWindow && canViewEntity;
 
   const {
     data: spendDataRaw,
@@ -150,7 +167,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
   } = usePaginatedDailyActivity({
     fetchFn: agentDailyActivityCall,
     args: [accessToken, startTime, endTime, null],
-    enabled: enabled && entityType === "team",
+    enabled: enabled && showAgentBreakdown,
   });
 
   const agentSpendData = agentSpendDataRaw as unknown as EntitySpendData;
@@ -158,7 +175,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
   const modelBreakdownKey = modelViewType === "groups" ? "model_groups" : "models";
   const modelMetrics = processActivityData(spendData, modelBreakdownKey, teams || []);
   const keyMetrics = processActivityData(spendData, "api_keys", teams || []);
-  const agentMetrics = entityType === "team" ? processActivityData(agentSpendData, "entities", teams || []) : {};
+  const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {};
 
   const getTopModels = () => {
     const modelSpend: { [key: string]: any } = {};
@@ -621,8 +638,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
         </Card>
       </Col>
 
-      {/* Top Agents - only for team entity type */}
-      {entityType === "team" && (
+      {showAgentBreakdown && (
         <Col numColSpan={2}>
           <Card>
             <Title>Top Agents Driving Spend</Title>
@@ -708,7 +724,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
         </>
       ),
     },
-    ...(entityType === "team"
+    ...(showAgentBreakdown
       ? [{ key: "agents", label: "Agent Activity", content: <ActivityMetrics modelMetrics={agentMetrics} /> }]
       : []),
     {
@@ -757,7 +773,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
           }
         />
       )}
-      {agentIsFetchingMore && entityType === "team" && (
+      {agentIsFetchingMore && showAgentBreakdown && (
         <Alert
           banner
           type="warning"
@@ -781,7 +797,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({ accessToken, entityType, enti
           }
         />
       )}
-      {agentCancelled && entityType === "team" && (
+      {agentCancelled && showAgentBreakdown && (
         <Alert
           banner
           type="info"
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
index 0ded7d195d0..841f0fda573 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
@@ -502,6 +502,8 @@ describe("UsagePage", () => {
     userId: "user-123",
     userEmail: "test@example.com",
     userRole: "Internal User",
+    userRoleLabel: "Internal User",
+    isViewOnly: false,
     premiumUser: true,
     disabledPersonalKeyCreation: false,
     showSSOBanner: false,
@@ -861,6 +863,29 @@ describe("UsagePage", () => {
     });
   });
 
+  // The select hides both views from a non-admin, so this drives the section
+  // gate directly through the mocked select, which always offers every option.
+  it.each(["organization", "agent"])("should not render the %s usage view for an internal user", async (usageView) => {
+    mockUseAuthorized.mockReturnValue(nonAdminSession);
+
+    renderWithProviders(<UsagePage {...defaultProps} organizations={mockOrganizations} />);
+
+    await waitFor(() => {
+      expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
+    });
+
+    const usageSelect = screen.getByTestId("usage-view-select");
+    act(() => {
+      fireEvent.change(usageSelect, { target: { value: "team" } });
+    });
+    expect(screen.getAllByText("Entity Usage").length).toBeGreaterThan(0);
+
+    act(() => {
+      fireEvent.change(usageSelect, { target: { value: usageView } });
+    });
+    expect(screen.queryByText("Entity Usage")).not.toBeInTheDocument();
+  });
+
   describe("admin user selector", () => {
     it("should render user selector for admin users in global view", async () => {
       renderWithProviders(<UsagePage {...defaultProps} />);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
index c3645d6371e..494df313ac0 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
@@ -33,6 +33,7 @@ import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
 import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers";
+import { hasCapability } from "@/utils/capabilities";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { all_admin_roles, internalUserRoles } from "@/utils/roles";
 import { ActivityMetrics, processActivityData } from "@/components/activity_metrics";
@@ -109,6 +110,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
   const { data: currentUser } = useCurrentUser();
   const isAdmin = all_admin_roles.includes(userRole || "");
   const canViewTagUsage = isAdmin || internalUserRoles.includes(userRole || "");
+  const canViewOrganizationUsage = hasCapability(userRole, "viewOrganizationUsage");
+  const canViewAgentUsage = hasCapability(userRole, "viewAgentUsage");
 
   // Debounced search for user selector
   const [userSearchInput, setUserSearchInput] = useState("");
@@ -513,7 +516,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
             <UsageViewSelect
               value={usageView}
               onChange={(value) => setUsageView(value)}
-              isAdmin={isAdmin}
+              userRole={userRole}
               canViewTagUsage={canViewTagUsage}
             />
             <AdvancedDatePicker value={dateValue} onValueChange={handleDateChange} />
@@ -950,7 +953,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
           )}
           {/* Organization Usage Panel */}
 
-          {usageView === "organization" && (
+          {usageView === "organization" && canViewOrganizationUsage && (
             <EntityUsage
               accessToken={accessToken}
               entityType="organization"
@@ -1033,7 +1036,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
               />
             </>
           )}
-          {usageView === "agent" && (
+          {usageView === "agent" && canViewAgentUsage && (
             <EntityUsage
               accessToken={accessToken}
               entityType="agent"
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
index aa9ff921ee1..9bc4bd81302 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
@@ -90,15 +90,16 @@ describe("UsageViewSelect", () => {
   });
 
   it("should render", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} isAdmin={false} />);
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" />);
 
     expect(screen.getByText("Usage View")).toBeInTheDocument();
     expect(screen.getByText("Select the usage data you want to view")).toBeInTheDocument();
     expect(screen.getByRole("combobox")).toBeInTheDocument();
+    expect(screen.getByRole("option", { name: "Your Usage" })).toBeInTheDocument();
   });
 
   it("should call onChange when value changes", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} isAdmin={true} />);
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Admin" />);
 
     const select = screen.getByRole("combobox");
     act(() => {
@@ -109,14 +110,34 @@ describe("UsageViewSelect", () => {
   });
 
   it("should show Tag Usage for non-admin users with tag usage permission", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} isAdmin={false} canViewTagUsage={true} />);
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
 
     expect(screen.getByRole("option", { name: "Tag Usage" })).toBeInTheDocument();
   });
 
   it("should hide Tag Usage for non-admin users without tag usage permission", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} isAdmin={false} />);
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" />);
 
     expect(screen.queryByRole("option", { name: "Tag Usage" })).not.toBeInTheDocument();
   });
+
+  it.each(["Organization Usage", "Agent Usage (A2A)"])("should show %s to an admin", (optionName) => {
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Admin" />);
+
+    expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument();
+  });
+
+  // Neither /organization/daily/activity nor /agent/daily/activity admits an
+  // internal user, so the option that fires them must not be selectable.
+  it.each(["Organization Usage", "Agent Usage (A2A)"])("should hide %s from an internal user", (optionName) => {
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
+
+    expect(screen.queryByRole("option", { name: optionName })).not.toBeInTheDocument();
+  });
+
+  it.each(["Team Usage", "Tag Usage"])("should keep %s available to an internal user", (optionName) => {
+    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
+
+    expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument();
+  });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx
index 94b483cb539..54c1d5ab7cc 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx
@@ -11,6 +11,8 @@ import {
 } from "@ant-design/icons";
 import { Badge, Select } from "antd";
 import React from "react";
+import { hasCapability, type Capability } from "@/utils/capabilities";
+import { all_admin_roles } from "@/utils/roles";
 export type UsageOption =
   | "global"
   | "my-usage"
@@ -24,7 +26,7 @@ export type UsageOption =
 export interface UsageViewSelectProps {
   value: UsageOption;
   onChange: (value: UsageOption) => void;
-  isAdmin: boolean;
+  userRole: string | null;
   canViewTagUsage?: boolean;
   title?: string;
   description?: string;
@@ -35,6 +37,7 @@ interface OptionConfig {
   label: string;
   description: string;
   icon: React.ReactNode;
+  capability?: Capability;
   adminOnly?: boolean;
   showForAdmin?: string;
   showForNonAdmin?: string;
@@ -63,12 +66,9 @@ const OPTIONS: OptionConfig[] = [
   {
     value: "organization",
     label: "Organization Usage",
-    showForAdmin: "Organization Usage",
-    showForNonAdmin: "Your Organization Usage",
-    description: "View organization-level usage",
-    descriptionForAdmin: "View usage across all organizations",
-    descriptionForNonAdmin: "View your organization's usage",
+    description: "View usage across all organizations",
     icon: <BankOutlined style={{ fontSize: "16px" }} />,
+    capability: "viewOrganizationUsage",
   },
   {
     value: "team",
@@ -95,7 +95,7 @@ const OPTIONS: OptionConfig[] = [
     label: "Agent Usage (A2A)",
     description: "View usage by AI agents",
     icon: <RobotOutlined style={{ fontSize: "16px" }} />,
-    adminOnly: true,
+    capability: "viewAgentUsage",
   },
   {
     value: "user",
@@ -115,14 +115,18 @@ const OPTIONS: OptionConfig[] = [
 export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
   value,
   onChange,
-  isAdmin,
+  userRole,
   canViewTagUsage = false,
   title = "Usage View",
   description = "Select the usage data you want to view",
   "data-id": dataId,
 }) => {
+  const isAdmin = all_admin_roles.includes(userRole ?? "");
   const getFilteredOptions = () => {
     return OPTIONS.filter((option) => {
+      if (option.capability) {
+        return hasCapability(userRole, option.capability);
+      }
       if (option.value === "tag" && canViewTagUsage) {
         return true;
       }
diff --git a/ui/litellm-dashboard/src/utils/capabilities.test.ts b/ui/litellm-dashboard/src/utils/capabilities.test.ts
index f48609b0b9d..3f5a8ac81fb 100644
--- a/ui/litellm-dashboard/src/utils/capabilities.test.ts
+++ b/ui/litellm-dashboard/src/utils/capabilities.test.ts
@@ -1,21 +1,31 @@
 import { describe, expect, it } from "vitest";
 
-import { hasCapability, rolesWithCapability } from "./capabilities";
+import { hasCapability, rolesWithCapability, type Capability } from "./capabilities";
+
+const ADMIN_ROLES = ["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer"];
+const NON_ADMIN_ROLES = [
+  "Internal User",
+  "Internal Viewer",
+  "App User",
+  "Org Admin",
+  "Unknown Role",
+  "",
+  null,
+  undefined,
+];
+
+const ADMIN_ONLY_CAPABILITIES: Capability[] = ["viewToolPolicies", "viewOrganizationUsage", "viewAgentUsage"];
 
 describe("hasCapability", () => {
-  it.each(["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer"])(
-    "should grant viewToolPolicies to %s",
-    (role) => {
-      expect(hasCapability(role, "viewToolPolicies")).toBe(true);
-    },
-  );
+  describe.each(ADMIN_ONLY_CAPABILITIES)("%s", (capability) => {
+    it.each(ADMIN_ROLES)("should grant it to %s", (role) => {
+      expect(hasCapability(role, capability)).toBe(true);
+    });
 
-  it.each(["Internal User", "Internal Viewer", "App User", "Org Admin", "Unknown Role", "", null, undefined])(
-    "should deny viewToolPolicies to %s",
-    (role) => {
-      expect(hasCapability(role, "viewToolPolicies")).toBe(false);
-    },
-  );
+    it.each(NON_ADMIN_ROLES)("should deny it to %s", (role) => {
+      expect(hasCapability(role, capability)).toBe(false);
+    });
+  });
 });
 
 describe("rolesWithCapability", () => {
diff --git a/ui/litellm-dashboard/src/utils/capabilities.ts b/ui/litellm-dashboard/src/utils/capabilities.ts
index 77ead2568fb..c4d878b81a5 100644
--- a/ui/litellm-dashboard/src/utils/capabilities.ts
+++ b/ui/litellm-dashboard/src/utils/capabilities.ts
@@ -2,6 +2,8 @@ import { all_admin_roles } from "./roles";
 
 const CAPABILITY_ROLES = {
   viewToolPolicies: all_admin_roles,
+  viewOrganizationUsage: all_admin_roles,
+  viewAgentUsage: all_admin_roles,
 } as const satisfies Record<string, readonly string[]>;
 
 export type Capability = keyof typeof CAPABILITY_ROLES;

From e3d3177ff1e1197cb3bf93d027b2463aa63b5bdf Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Sat, 8 Aug 2026 21:12:27 -0700
Subject: [PATCH 2/6] style(ui): drop narration comments from the usage gating
 tests

Both restated what the test name and the surrounding setup already say, so
they were maintenance cost without explanatory value. The reasoning they
carried lives in the commit that added the gates.
---
 .../usage/_components/components/UsagePageView.test.tsx         | 2 --
 .../components/UsageViewSelect/UsageViewSelect.test.tsx         | 2 --
 2 files changed, 4 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
index 841f0fda573..9085cf961a9 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
@@ -863,8 +863,6 @@ describe("UsagePage", () => {
     });
   });
 
-  // The select hides both views from a non-admin, so this drives the section
-  // gate directly through the mocked select, which always offers every option.
   it.each(["organization", "agent"])("should not render the %s usage view for an internal user", async (usageView) => {
     mockUseAuthorized.mockReturnValue(nonAdminSession);
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
index 9bc4bd81302..dcc0ce06673 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
@@ -127,8 +127,6 @@ describe("UsageViewSelect", () => {
     expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument();
   });
 
-  // Neither /organization/daily/activity nor /agent/daily/activity admits an
-  // internal user, so the option that fires them must not be selectable.
   it.each(["Organization Usage", "Agent Usage (A2A)"])("should hide %s from an internal user", (optionName) => {
     render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
 

From 00da19e4e8afa8119a97ad93a8094ae74884b391 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Mon, 10 Aug 2026 14:31:44 -0700
Subject: [PATCH 3/6] refactor(ui): extract entity usage aggregations into
 their own module

Merging staging's flat-cost summary work with the capability gating pushed
EntityUsage.tsx to 815 counted lines, over the 800-line eslint cap. Move the
four pure top-N/rollup helpers to entityUsageAggregations.ts and pass their
inputs explicitly.

TopKeyView and TopModelView were mocked to render static text, so nothing
asserted which breakdown fed which table. The mocks now surface their rows and
a new case pins each table to its own data source.
---
 .../EntityUsage/EntityUsage.test.tsx          |  49 ++++-
 .../components/EntityUsage/EntityUsage.tsx    | 186 ++----------------
 .../EntityUsage/entityUsageAggregations.ts    | 168 ++++++++++++++++
 3 files changed, 228 insertions(+), 175 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
index 11528117f1e..c85a9fb71f6 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
@@ -39,11 +39,21 @@ vi.mock("../EndpointUsage/EndpointUsage", () => ({
 }));
 
 vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({
-  default: () => <div>Top Keys</div>,
+  default: ({ topKeys }: { topKeys: { api_key: string; spend: number }[] }) => (
+    <div>
+      <span>Top Keys</span>
+      <span>{`top-keys:${topKeys.map((row) => `${row.api_key}=${row.spend}`).join("|")}`}</span>
+    </div>
+  ),
 }));
 
 vi.mock("./TopModelView", () => ({
-  default: () => <div>Top Models</div>,
+  default: ({ topModels }: { topModels: { key: string; spend: number }[] }) => (
+    <div>
+      <span>Top Models</span>
+      <span>{`top-models:${topModels.map((row) => `${row.key}=${row.spend}`).join("|")}`}</span>
+    </div>
+  ),
 }));
 
 vi.mock("@/components/EntityUsageExport/EntityUsageExportModal", () => ({
@@ -922,4 +932,39 @@ describe("EntityUsage", () => {
     expect(screen.queryByAltText("zzz-internal logo")).not.toBeInTheDocument();
     expect(screen.getByText("z")).toBeInTheDocument();
   });
+
+  it("feeds the key, model and agent tables from their own breakdowns", async () => {
+    const usageMetrics = {
+      spend: 30.75,
+      api_requests: 300,
+      successful_requests: 290,
+      failed_requests: 10,
+      total_tokens: 15000,
+      prompt_tokens: 9000,
+      completion_tokens: 6000,
+      cache_read_input_tokens: 0,
+      cache_creation_input_tokens: 0,
+    };
+    mockTeamDailyActivityCall.mockResolvedValue({
+      ...mockSpendData,
+      results: [
+        {
+          ...mockSpendData.results[0],
+          breakdown: {
+            ...mockSpendData.results[0].breakdown,
+            model_groups: { "gpt-4o": { metrics: { ...usageMetrics, spend: 70.25 }, metadata: {} } },
+            api_keys: { "sk-abc": { metrics: usageMetrics, metadata: { key_alias: "prod-key", team_id: null } } },
+          },
+        },
+      ],
+    });
+
+    render(<EntityUsage {...defaultProps} entityType="team" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("top-keys:sk-abc=30.75")).toBeInTheDocument();
+    });
+    expect(screen.getByText("top-models:gpt-4o=70.25")).toBeInTheDocument();
+    expect(screen.getByText(/^top-models:Code Review Agent=/)).toBeInTheDocument();
+  });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
index 0d1ac58397b..956060fc244 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
@@ -1,5 +1,12 @@
 import useTeams from "@/app/(dashboard)/hooks/useTeams";
 import { BarChart, DonutChart } from "@/components/shared/charts";
+import {
+  getProviderSpend,
+  getTopAgents,
+  getTopAPIKeys,
+  getTopModels,
+  type ExtendedDailyData,
+} from "./entityUsageAggregations";
 import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost, type SummaryTile } from "./entityUsageSummary";
 import { MoneyCell } from "@/components/shared/table_cells";
 import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
@@ -42,13 +49,7 @@ import {
 } from "@/components/networking";
 import { Logo } from "@/components/molecules/logo/Logo";
 import { usePaginatedDailyActivity } from "../../hooks/usePaginatedDailyActivity";
-import {
-  BreakdownMetrics,
-  DailyData,
-  EntityMetricWithMetadata,
-  KeyMetricWithMetadata,
-  TagUsage,
-} from "@/components/UsagePage/types";
+import { EntityMetricWithMetadata } from "@/components/UsagePage/types";
 import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatters";
 import EndpointUsage from "../EndpointUsage/EndpointUsage";
 import ModelViewToggle, { ModelViewType } from "../ModelViewToggle";
@@ -70,10 +71,6 @@ interface EntityMetrics {
   metadata: Record<string, any>;
 }
 
-type ExtendedDailyData = DailyData & {
-  breakdown: BreakdownMetrics;
-};
-
 interface EntitySpendData {
   results: ExtendedDailyData[];
   metadata: {
@@ -180,163 +177,6 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
   const keyMetrics = processActivityData(spendData, "api_keys", teams || []);
   const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {};
 
-  const getTopModels = () => {
-    const modelSpend: { [key: string]: any } = {};
-    spendData.results.forEach((day) => {
-      Object.entries(day.breakdown[modelBreakdownKey] || {}).forEach(([model, metrics]) => {
-        if (!modelSpend[model]) {
-          modelSpend[model] = {
-            spend: 0,
-            requests: 0,
-            successful_requests: 0,
-            failed_requests: 0,
-            tokens: 0,
-          };
-        }
-        try {
-          modelSpend[model].spend += metrics.metrics.spend;
-        } catch (e) {
-          console.error(`Error adding spend for ${model}: ${e}, got metrics: ${JSON.stringify(metrics)}`);
-        }
-        modelSpend[model].requests += metrics.metrics.api_requests;
-        modelSpend[model].successful_requests += metrics.metrics.successful_requests;
-        modelSpend[model].failed_requests += metrics.metrics.failed_requests;
-        modelSpend[model].tokens += metrics.metrics.total_tokens;
-      });
-    });
-
-    return Object.entries(modelSpend)
-      .map(([model, metrics]) => ({
-        key: model,
-        ...metrics,
-      }))
-      .sort((a, b) => b.spend - a.spend)
-      .slice(0, topModelsLimit);
-  };
-
-  const getTopAgents = () => {
-    const agentSpend: { [key: string]: any } = {};
-    agentSpendData.results.forEach((day) => {
-      Object.entries(day.breakdown.entities || {}).forEach(([agentId, data]) => {
-        if (!agentSpend[agentId]) {
-          agentSpend[agentId] = {
-            spend: 0,
-            requests: 0,
-            successful_requests: 0,
-            failed_requests: 0,
-            tokens: 0,
-            agent_name: (data.metadata as any)?.agent_name || agentId,
-          };
-        }
-        agentSpend[agentId].spend += data.metrics.spend;
-        agentSpend[agentId].requests += data.metrics.api_requests;
-        agentSpend[agentId].successful_requests += data.metrics.successful_requests;
-        agentSpend[agentId].failed_requests += data.metrics.failed_requests;
-        agentSpend[agentId].tokens += data.metrics.total_tokens;
-      });
-    });
-
-    return Object.entries(agentSpend)
-      .map(([agentId, metrics]) => ({
-        key: metrics.agent_name,
-        ...metrics,
-      }))
-      .sort((a, b) => b.spend - a.spend)
-      .slice(0, topAgentsLimit);
-  };
-
-  const getTopAPIKeys = () => {
-    const keySpend: { [key: string]: KeyMetricWithMetadata } = {};
-    spendData.results.forEach((day) => {
-      const { breakdown } = day;
-      const { entities } = breakdown;
-      const tagDictionary = Object.keys(entities).reduce((acc: { [key: string]: TagUsage[] }, entity) => {
-        const { api_key_breakdown } = entities[entity];
-        Object.keys(api_key_breakdown).forEach((key) => {
-          const tagUsage = { tag: entity, usage: api_key_breakdown[key].metrics.spend };
-          if (acc[key]) {
-            acc[key].push(tagUsage);
-          } else {
-            acc[key] = [tagUsage];
-          }
-        });
-        return acc;
-      }, {});
-      Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => {
-        if (!keySpend[key]) {
-          keySpend[key] = {
-            metrics: {
-              spend: 0,
-              prompt_tokens: 0,
-              completion_tokens: 0,
-              total_tokens: 0,
-              api_requests: 0,
-              successful_requests: 0,
-              failed_requests: 0,
-              cache_read_input_tokens: 0,
-              cache_creation_input_tokens: 0,
-            },
-            metadata: {
-              key_alias: metrics.metadata.key_alias,
-              team_id: metrics.metadata.team_id || null,
-              tags: tagDictionary[key] || [],
-            },
-          };
-        }
-        keySpend[key].metrics.spend += metrics.metrics.spend;
-        keySpend[key].metrics.prompt_tokens += metrics.metrics.prompt_tokens;
-        keySpend[key].metrics.completion_tokens += metrics.metrics.completion_tokens;
-        keySpend[key].metrics.total_tokens += metrics.metrics.total_tokens;
-        keySpend[key].metrics.api_requests += metrics.metrics.api_requests;
-        keySpend[key].metrics.successful_requests += metrics.metrics.successful_requests;
-        keySpend[key].metrics.failed_requests += metrics.metrics.failed_requests;
-        keySpend[key].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0;
-        keySpend[key].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0;
-      });
-    });
-
-    return Object.entries(keySpend)
-      .map(([api_key, metrics]) => ({
-        api_key,
-        key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias
-        tags: metrics.metadata.tags || "-",
-        spend: metrics.metrics.spend,
-      }))
-      .sort((a, b) => b.spend - a.spend)
-      .slice(0, topKeysLimit);
-  };
-
-  const getProviderSpend = () => {
-    const providerSpend: { [key: string]: any } = {};
-    spendData.results.forEach((day) => {
-      Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => {
-        if (!providerSpend[provider]) {
-          providerSpend[provider] = {
-            provider,
-            spend: 0,
-            requests: 0,
-            successful_requests: 0,
-            failed_requests: 0,
-            tokens: 0,
-          };
-        }
-        try {
-          providerSpend[provider].spend += metrics.metrics.spend;
-          providerSpend[provider].requests += metrics.metrics.api_requests;
-          providerSpend[provider].successful_requests += metrics.metrics.successful_requests;
-          providerSpend[provider].failed_requests += metrics.metrics.failed_requests;
-          providerSpend[provider].tokens += metrics.metrics.total_tokens;
-        } catch (e) {
-          console.error(`Error processing provider ${provider}: ${e}`);
-        }
-      });
-    });
-
-    return Object.values(providerSpend)
-      .filter((provider) => provider.spend > 0)
-      .sort((a, b) => b.spend - a.spend);
-  };
-
   const getAllTags = () => {
     if (entityList) {
       return entityList;
@@ -633,7 +473,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
         <Card>
           <Title>Top Virtual Keys</Title>
           <TopKeyView
-            topKeys={getTopAPIKeys()}
+            topKeys={getTopAPIKeys(spendData.results, topKeysLimit)}
             teams={null}
             showTags={entityType === "tag"}
             topKeysLimit={topKeysLimit}
@@ -650,7 +490,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
             <ModelViewToggle value={modelViewType} onChange={setModelViewType} />
           </div>
           <TopModelView
-            topModels={getTopModels()}
+            topModels={getTopModels(spendData.results, modelBreakdownKey, topModelsLimit)}
             topModelsLimit={topModelsLimit}
             setTopModelsLimit={setTopModelsLimit}
           />
@@ -662,7 +502,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
           <Card>
             <Title>Top Agents Driving Spend</Title>
             <TopModelView
-              topModels={getTopAgents()}
+              topModels={getTopAgents(agentSpendData.results, topAgentsLimit)}
               topModelsLimit={topAgentsLimit}
               setTopModelsLimit={setTopAgentsLimit}
             />
@@ -679,7 +519,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
               <Col numColSpan={1}>
                 <DonutChart
                   className="mt-4 h-40"
-                  data={getProviderSpend()}
+                  data={getProviderSpend(spendData.results)}
                   index="provider"
                   category="spend"
                   valueFormatter={(value) => `$${formatNumberWithCommas(value, 2)}`}
@@ -701,7 +541,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
                     </TableRow>
                   </TableHead>
                   <TableBody>
-                    {getProviderSpend().map((provider) => (
+                    {getProviderSpend(spendData.results).map((provider) => (
                       <TableRow key={provider.provider}>
                         <TableCell>
                           <div className="flex items-center space-x-2">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts
new file mode 100644
index 00000000000..fc72b66f974
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageAggregations.ts
@@ -0,0 +1,168 @@
+import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types";
+
+export type ExtendedDailyData = DailyData & {
+  breakdown: BreakdownMetrics;
+};
+
+export type ModelBreakdownKey = "models" | "model_groups";
+
+export const getTopModels = (
+  results: ExtendedDailyData[],
+  modelBreakdownKey: ModelBreakdownKey,
+  topModelsLimit: number,
+) => {
+  const modelSpend: { [key: string]: any } = {};
+  results.forEach((day) => {
+    Object.entries(day.breakdown[modelBreakdownKey] || {}).forEach(([model, metrics]) => {
+      if (!modelSpend[model]) {
+        modelSpend[model] = {
+          spend: 0,
+          requests: 0,
+          successful_requests: 0,
+          failed_requests: 0,
+          tokens: 0,
+        };
+      }
+      try {
+        modelSpend[model].spend += metrics.metrics.spend;
+      } catch (e) {
+        console.error(`Error adding spend for ${model}: ${e}, got metrics: ${JSON.stringify(metrics)}`);
+      }
+      modelSpend[model].requests += metrics.metrics.api_requests;
+      modelSpend[model].successful_requests += metrics.metrics.successful_requests;
+      modelSpend[model].failed_requests += metrics.metrics.failed_requests;
+      modelSpend[model].tokens += metrics.metrics.total_tokens;
+    });
+  });
+
+  return Object.entries(modelSpend)
+    .map(([model, metrics]) => ({
+      key: model,
+      ...metrics,
+    }))
+    .sort((a, b) => b.spend - a.spend)
+    .slice(0, topModelsLimit);
+};
+
+export const getTopAgents = (results: ExtendedDailyData[], topAgentsLimit: number) => {
+  const agentSpend: { [key: string]: any } = {};
+  results.forEach((day) => {
+    Object.entries(day.breakdown.entities || {}).forEach(([agentId, data]) => {
+      if (!agentSpend[agentId]) {
+        agentSpend[agentId] = {
+          spend: 0,
+          requests: 0,
+          successful_requests: 0,
+          failed_requests: 0,
+          tokens: 0,
+          agent_name: (data.metadata as any)?.agent_name || agentId,
+        };
+      }
+      agentSpend[agentId].spend += data.metrics.spend;
+      agentSpend[agentId].requests += data.metrics.api_requests;
+      agentSpend[agentId].successful_requests += data.metrics.successful_requests;
+      agentSpend[agentId].failed_requests += data.metrics.failed_requests;
+      agentSpend[agentId].tokens += data.metrics.total_tokens;
+    });
+  });
+
+  return Object.entries(agentSpend)
+    .map(([agentId, metrics]) => ({
+      key: metrics.agent_name,
+      ...metrics,
+    }))
+    .sort((a, b) => b.spend - a.spend)
+    .slice(0, topAgentsLimit);
+};
+
+export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number) => {
+  const keySpend: { [key: string]: KeyMetricWithMetadata } = {};
+  results.forEach((day) => {
+    const { breakdown } = day;
+    const { entities } = breakdown;
+    const tagDictionary = Object.keys(entities).reduce((acc: { [key: string]: TagUsage[] }, entity) => {
+      const { api_key_breakdown } = entities[entity];
+      Object.keys(api_key_breakdown).forEach((key) => {
+        const tagUsage = { tag: entity, usage: api_key_breakdown[key].metrics.spend };
+        if (acc[key]) {
+          acc[key].push(tagUsage);
+        } else {
+          acc[key] = [tagUsage];
+        }
+      });
+      return acc;
+    }, {});
+    Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => {
+      if (!keySpend[key]) {
+        keySpend[key] = {
+          metrics: {
+            spend: 0,
+            prompt_tokens: 0,
+            completion_tokens: 0,
+            total_tokens: 0,
+            api_requests: 0,
+            successful_requests: 0,
+            failed_requests: 0,
+            cache_read_input_tokens: 0,
+            cache_creation_input_tokens: 0,
+          },
+          metadata: {
+            key_alias: metrics.metadata.key_alias,
+            team_id: metrics.metadata.team_id || null,
+            tags: tagDictionary[key] || [],
+          },
+        };
+      }
+      keySpend[key].metrics.spend += metrics.metrics.spend;
+      keySpend[key].metrics.prompt_tokens += metrics.metrics.prompt_tokens;
+      keySpend[key].metrics.completion_tokens += metrics.metrics.completion_tokens;
+      keySpend[key].metrics.total_tokens += metrics.metrics.total_tokens;
+      keySpend[key].metrics.api_requests += metrics.metrics.api_requests;
+      keySpend[key].metrics.successful_requests += metrics.metrics.successful_requests;
+      keySpend[key].metrics.failed_requests += metrics.metrics.failed_requests;
+      keySpend[key].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0;
+      keySpend[key].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0;
+    });
+  });
+
+  return Object.entries(keySpend)
+    .map(([api_key, metrics]) => ({
+      api_key,
+      key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias
+      tags: metrics.metadata.tags || "-",
+      spend: metrics.metrics.spend,
+    }))
+    .sort((a, b) => b.spend - a.spend)
+    .slice(0, topKeysLimit);
+};
+
+export const getProviderSpend = (results: ExtendedDailyData[]) => {
+  const providerSpend: { [key: string]: any } = {};
+  results.forEach((day) => {
+    Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => {
+      if (!providerSpend[provider]) {
+        providerSpend[provider] = {
+          provider,
+          spend: 0,
+          requests: 0,
+          successful_requests: 0,
+          failed_requests: 0,
+          tokens: 0,
+        };
+      }
+      try {
+        providerSpend[provider].spend += metrics.metrics.spend;
+        providerSpend[provider].requests += metrics.metrics.api_requests;
+        providerSpend[provider].successful_requests += metrics.metrics.successful_requests;
+        providerSpend[provider].failed_requests += metrics.metrics.failed_requests;
+        providerSpend[provider].tokens += metrics.metrics.total_tokens;
+      } catch (e) {
+        console.error(`Error processing provider ${provider}: ${e}`);
+      }
+    });
+  });
+
+  return Object.values(providerSpend)
+    .filter((provider) => provider.spend > 0)
+    .sort((a, b) => b.spend - a.spend);
+};

From ade5a425e8bb3ab60858d35f2c473a0b1d830b93 Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Mon, 10 Aug 2026 14:37:09 -0700
Subject: [PATCH 4/6] fix(proxy): isolate guardrail load failures per row
 (#36432)

* fix(proxy): isolate guardrail load failures per row

One DB guardrail row that fails to initialize aborted the whole
_init_guardrails_in_db loop, so a single typo'd guardrail type or a
missing required param left the proxy running with zero DB guardrails
registered and requests that should have been blocked reaching the
provider.

Catch per row around sync_guardrail_from_db, log the guardrail name, id
and error, and continue with the remaining rows. The failing row's id is
still added to db_guardrail_ids before the attempt so reconcile_db_guardrails
cannot mistake a live row for a deleted one.

* test(proxy): drop inline note and record reconcile via a handler double

Replaces the patched bound method with an InMemoryGuardrailHandler subclass
that records what reconcile_db_guardrails received, so the test injects a
double instead of swapping a method on a live object.
---
 litellm/proxy/proxy_server.py                 | 16 ++++-
 .../proxy/proxy_server/test_proxy_config.py   | 64 +++++++++++++++++++
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 3a4896dca9e..bc980934f9f 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -6860,9 +6860,19 @@ class ProxyConfig:
                 guardrail_id = guardrail.get("guardrail_id")
                 if guardrail_id:
                     db_guardrail_ids.add(guardrail_id)
-                IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
-                    guardrail=cast(Guardrail, guardrail),
-                )
+                try:
+                    IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
+                        guardrail=cast(Guardrail, guardrail),
+                    )
+                except Exception as e:  # noqa: BLE001  # one unloadable row must not stop the remaining guardrails
+                    verbose_proxy_logger.error(
+                        "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - "
+                        "skipping guardrail '%s' (ID: %s): %s: %s",
+                        guardrail.get("guardrail_name"),
+                        guardrail_id,
+                        type(e).__name__,
+                        e,
+                    )
 
             # Drop in-memory DB-backed entries whose row was deleted on another
             # pod. Config-loaded entries are never touched.
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 91a7e1bc2c2..f70be17eb95 100644
--- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
+++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py
@@ -2639,3 +2639,67 @@ async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembe
     assert clean_agent_registry.config_agents == ()
     clean_agent_registry.load_agents_from_db_and_config(db_agents=None)
     assert clean_agent_registry.get_agent_list() == ()
+
+
+# ---------------------------------------------------------------------------
+# _init_guardrails_in_db
+# ---------------------------------------------------------------------------
+
+
+def _db_guardrail_row(guardrail_id: str, guardrail_type: str) -> dict[str, object]:
+    return {
+        "guardrail_id": guardrail_id,
+        "guardrail_name": f"name-{guardrail_id}",
+        "litellm_params": {"guardrail": guardrail_type, "mode": "pre_call"},
+        "guardrail_info": None,
+        "team_id": None,
+    }
+
+
+@pytest.mark.asyncio
+async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(monkeypatch):
+    """
+    A single DB row that fails to initialize used to abort the whole loop, so one
+    typo'd guardrail type left the proxy running with zero guardrails loaded.
+
+    The failing row's id must still reach reconcile_db_guardrails so that eviction
+    pass cannot treat a row that is alive in the DB as one that was deleted.
+    """
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.proxy.guardrails import guardrail_registry as registry_module
+    from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams
+
+    class _RecordingHandler(registry_module.InMemoryGuardrailHandler):
+        def __init__(self) -> None:
+            super().__init__()
+            self.reconciled_with: list[set[str]] = []
+
+        def reconcile_db_guardrails(self, db_guardrail_ids: set[str]) -> list[str]:
+            self.reconciled_with.append(set(db_guardrail_ids))
+            return super().reconcile_db_guardrails(db_guardrail_ids)
+
+    handler = _RecordingHandler()
+    monkeypatch.setattr(registry_module, "IN_MEMORY_GUARDRAIL_HANDLER", handler)
+
+    def _initializer(litellm_params: LitellmParams, guardrail: Guardrail) -> CustomGuardrail:
+        return CustomGuardrail(
+            guardrail_name=guardrail["guardrail_name"],
+            event_hook=GuardrailEventHooks.pre_call,
+            default_on=False,
+        )
+
+    monkeypatch.setitem(registry_module.guardrail_initializer_registry, "lit5367_ok", _initializer)
+
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
+        return_value=[
+            _db_guardrail_row("first", "lit5367_ok"),
+            _db_guardrail_row("broken", "litellm_tool_permission"),
+            _db_guardrail_row("last", "lit5367_ok"),
+        ]
+    )
+
+    await ProxyConfig()._init_guardrails_in_db(prisma_client=prisma_client)
+
+    assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"]
+    assert handler.reconciled_with == [{"first", "broken", "last"}]

From c40828509b7c73399c31556b9a1e37dda2f202b5 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Mon, 10 Aug 2026 14:42:36 -0700
Subject: [PATCH 5/6] fix(reset_budget_job): atomic budget cascade with chunked
 reset scans (#36287)

* fix(reset_budget_job): advance budget_reset_at atomically with the spend cascade

A postgres timeout mid-cascade previously left LiteLLM_BudgetTable rows
stamped for the next window while team member, enduser, org and tag spend
stayed at cap, so every later tick skipped them until the window rolled
over. All cascade writes and the budget_reset_at advance now share one
prisma batch transaction; a failed run persists nothing and the rows stay
due for the next ~10 minute tick. Cache and counter invalidation runs only
after commit, and the catch-all enduser log line now names the cascade.

* fix(reset_budget_job): elect one runner per tick and chunk the reset scans

Every pod and worker previously ran the reset job every ~10 minutes,
each fetching every expired row with no limit and writing one giant
transaction at the same calendar-aligned boundary; that concurrency is
what piled up postgres lock contention and timeouts. The job now takes
the shared PodLockManager redis lock (no redis keeps the old behavior),
and each phase walks its due rows in 500-row chunks, one transaction per
chunk, stopping when a chunk is short, makes no forward progress, or
hits the per-run cap; leftovers wait for the next tick.

* chore(lint): ratchet budget ceilings down for fixed violations

* fix(reset_budget_job): harden chunk loop, fail open on redis errors, heartbeat the lock

Review fixes on the two prior commits. Reset scans now skip rows with no
budget_duration, so permanently due rows can neither starve a phase nor
have a lifetime cap zeroed every tick. Chunk progress counts rows whose
new budget_reset_at actually cleared the cutoff, so a zero-length
duration cannot burn the per-run chunk cap. A failed lock acquire only
skips the run when another pod verifiably holds the lock; a broken redis
runs unguarded instead of silently disabling resets fleet-wide. Partial
row failures report real progress and fire the failure hook without
killing the phase. The leader re-asserts the lock between phases and
stops if another pod took over, and the budget window advance uses
update_many so a tier deleted mid-chunk cannot abort the transaction.
Lint budget ceilings re-ratcheted for the net-fixed violations.

* fix(reset_budget_job): renew the leader lease and reject non-positive budget durations

Bot review follow-ups. PodLockManager now extends the lock TTL when the
holding pod re-acquires, via an atomic compare-and-expire script with a
plain SET fallback, so a run longer than the TTL keeps its lease instead
of silently sharing the job with another pod. The positive-duration
validation that team member endpoints already had is hoisted to
management common_utils and applied to key, internal user, budget,
customer and team intake, so a tenant can no longer create zero-duration
budgets whose permanently due rows starve other tenants' resets. Such
durations now return 400 at intake; existing rows are untouched.

* refactor(reset_budget_job): defer leader election to a follow-up PR

* fix(reset_budget_job): satisfy strict lint gates

String defaults for the two getenv calls (PLW1508) and the chunk
outcome returns moved to try/else (TRY300).
---
 basedpyright-code-budget.json                 |    6 +-
 litellm/constants.py                          |    2 +
 .../proxy/common_utils/reset_budget_job.py    |  705 ++++----
 .../budget_management_endpoints.py            |    9 +-
 .../management_endpoints/common_utils.py      |   29 +
 .../customer_endpoints.py                     |    2 +
 .../internal_user_endpoints.py                |    4 +
 .../key_management_endpoints.py               |    6 +-
 .../management_endpoints/team_endpoints.py    |   34 +-
 litellm/proxy/utils.py                        |   39 +-
 litellm/repositories/__init__.py              |    8 +
 litellm/repositories/prisma_protocols.py      |   17 +
 litellm/repositories/unit_of_work.py          |   63 +-
 ruff-strict-budget.json                       |    8 +-
 .../test_proxy_budget_reset.py                |  442 +++--
 .../common_utils/test_reset_budget_job.py     | 1479 +++++++++--------
 .../test_budget_endpoints.py                  |   33 +
 .../management_endpoints/test_common_utils.py |   52 +
 .../test_customer_endpoints.py                |   34 +
 .../test_internal_user_endpoints.py           |   62 +
 .../test_key_management_endpoints.py          |   68 +-
 .../test_team_endpoints.py                    |   60 +
 .../repositories/test_unit_of_work.py         |   63 +-
 type-discipline-budget.json                   |    8 +-
 24 files changed, 1937 insertions(+), 1296 deletions(-)

diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json
index 65d3c239253..96b689aed74 100644
--- a/basedpyright-code-budget.json
+++ b/basedpyright-code-budget.json
@@ -99,19 +99,19 @@
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 45004
+    "limit": 44996
   },
   "reportUnknownLambdaType": {
     "limit": 113
   },
   "reportUnknownMemberType": {
-    "limit": 39649
+    "limit": 39643
   },
   "reportUnknownParameterType": {
     "limit": 20132
   },
   "reportUnknownVariableType": {
-    "limit": 31156
+    "limit": 31153
   },
   "reportUnnecessaryCast": {
     "limit": 118
diff --git a/litellm/constants.py b/litellm/constants.py
index f2ac96162eb..87d6fa1a744 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1493,6 +1493,8 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT
 SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
 DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60))  # 1 minute
 PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
+RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
+RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100")))
 PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
 MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
 MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py
index 8830970f96f..bf760a92d88 100644
--- a/litellm/proxy/common_utils/reset_budget_job.py
+++ b/litellm/proxy/common_utils/reset_budget_job.py
@@ -1,14 +1,21 @@
 import asyncio
 import json
 import time
-from collections.abc import Callable, Sequence
+from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
+from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import Final, Literal, Protocol, TypeVar
+from types import MappingProxyType
+from typing import Final, Literal, Protocol, TypeVar, assert_never
 
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.dual_cache import DualCache
-from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME
+from litellm.constants import (
+    GLOBAL_PROXY_SPEND_CACHE_KEY,
+    LITELLM_PROXY_BUDGET_NAME,
+    RESET_BUDGET_JOB_BATCH_SIZE,
+    RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN,
+)
 from litellm.proxy._types import (
     LiteLLM_BudgetTableFull,
     LiteLLM_EndUserTable,
@@ -30,7 +37,10 @@ from litellm.repositories.table_repositories import (
     TeamMembershipRepository,
 )
 from litellm.repositories.team_repository import TeamRepository
-from litellm.repositories.unit_of_work import spend_reset_unit_of_work
+from litellm.repositories.unit_of_work import (
+    budget_cascade_unit_of_work,
+    spend_reset_unit_of_work,
+)
 from litellm.repositories.verification_token_repository import (
     VerificationTokenRepository,
 )
@@ -38,6 +48,9 @@ from litellm.types.services import ServiceTypes
 
 _RowT = TypeVar("_RowT")
 
+_LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}})
+_SPENT_ROWS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"spend": {"gt": 0}})
+
 
 class _TeamMembershipRow(Protocol):
     @property
@@ -62,39 +75,130 @@ class _TagRow(Protocol):
     def tag_name(self) -> str: ...
 
 
+class _EndUserRow(Protocol):
+    @property
+    def user_id(self) -> str: ...
+
+
 def _team_membership_counter_key(row: _TeamMembershipRow) -> str:
     return f"spend:team_member:{row.user_id}:{row.team_id}"
 
 
-def _team_membership_cache_key(row: _TeamMembershipRow) -> str:
-    return f"{row.team_id}_{row.user_id}"
+def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]:
+    return (f"{row.team_id}_{row.user_id}",)
 
 
 def _key_counter_key(row: _KeyRow) -> str:
     return f"spend:key:{row.token}"
 
 
-def _key_cache_key(row: _KeyRow) -> str:
-    return row.token
+def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]:
+    return (row.token,)
 
 
 def _org_counter_key(row: _OrgRow) -> str:
     return f"spend:org:{row.organization_id}"
 
 
-def _org_cache_keys(row: _OrgRow) -> Sequence[str]:
-    return [
+def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]:
+    return (
         f"org_id:{row.organization_id}",
         f"org_id:{row.organization_id}:with_budget",
-    ]
+    )
 
 
 def _tag_counter_key(row: _TagRow) -> str:
     return f"spend:tag:{row.tag_name}"
 
 
-def _tag_cache_key(row: _TagRow) -> str:
-    return f"tag:{row.tag_name}"
+def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]:
+    return (f"tag:{row.tag_name}",)
+
+
+def _budget_link_where(
+    budget_ids: Sequence[str],
+    extra: Mapping[str, object] = MappingProxyType({}),
+) -> dict[str, object]:
+    return {"budget_id": {"in": list(budget_ids)}, **extra}
+
+
+@dataclass(frozen=True, slots=True)
+class _BudgetCascade:
+    """Everything one budget-tier reset touches, resolved before any write."""
+
+    budgets: tuple[LiteLLM_BudgetTableFull, ...] = ()
+    budget_ids: tuple[str, ...] = ()
+    budget_resets: tuple[tuple[str, datetime], ...] = ()
+    endusers: tuple[_EndUserRow, ...] = ()
+    counter_keys: tuple[str, ...] = ()
+    cache_keys: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class _BudgetCascadeCommitted:
+    cascade: _BudgetCascade
+    advanced: int
+
+
+@dataclass(frozen=True, slots=True)
+class _BudgetCascadeFailed:
+    cascade: _BudgetCascade
+    error: Exception
+
+
+_EMPTY_CASCADE: Final = _BudgetCascade()
+
+
+@dataclass(frozen=True, slots=True)
+class _ChunkOutcome:
+    """One chunk of a reset phase: rows read, and rows whose new budget_reset_at
+    cleared the due cutoff. Anything else is still due and would come straight
+    back on the next fetch, so it is not progress."""
+
+    fetched: int
+    advanced: int
+
+
+_NO_PROGRESS: Final = _ChunkOutcome(fetched=0, advanced=0)
+
+
+def _as_utc(moment: datetime) -> datetime:
+    return moment if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc)
+
+
+def _count_advanced(reset_ats: Iterable[object], cutoff: datetime) -> int:
+    """How many rows the write actually moved past the due cutoff.
+
+    A budget_duration of "0s" (or one the parser cannot read) resolves to the
+    current time, so the row is written and stays due. Counting it as progress
+    would re-read the same chunk until the per-run cap on every tick.
+    """
+    utc_cutoff: Final = _as_utc(cutoff)
+    return sum(1 for reset_at in reset_ats if isinstance(reset_at, datetime) and _as_utc(reset_at) > utc_cutoff)
+
+
+def _phase_is_drained(outcome: _ChunkOutcome) -> bool:
+    """A short chunk means the due rows ran out. A full chunk that advanced
+    nothing would be re-read unchanged forever, so it ends the phase too and
+    those rows wait for the next tick."""
+    return outcome.fetched < RESET_BUDGET_JOB_BATCH_SIZE or outcome.advanced == 0
+
+
+async def _run_phase_in_chunks(process_chunk: Callable[[], Awaitable[_ChunkOutcome]]) -> None:
+    """Drive one reset phase a chunk at a time, capped so a single run cannot
+    spin unbounded: leftovers are picked up by the next tick."""
+    for _ in range(RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN):
+        if _phase_is_drained(await process_chunk()):
+            return
+
+
+def _budget_cascade_event_metadata(cascade: _BudgetCascade) -> dict[str, object]:
+    return {
+        "num_budgets_found": len(cascade.budgets),
+        "budgets_found": json.dumps(cascade.budgets, indent=4, default=str),
+        "num_endusers_found": len(cascade.endusers),
+        "endusers_found": json.dumps(cascade.endusers, indent=4, default=str),
+    }
 
 
 class ResetBudgetJob:
@@ -122,21 +226,14 @@ class ResetBudgetJob:
 
         Updates db
         """
-        if self.prisma_client is not None:
-            ### RESET KEY BUDGET ###
-            await self.reset_budget_for_litellm_keys()
+        if self.prisma_client is None:
+            return
 
-            ### RESET USER BUDGET ###
-            await self.reset_budget_for_litellm_users()
-
-            ## Reset Team Budget
-            await self.reset_budget_for_litellm_teams()
-
-            ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ###
-            await self.reset_budget_for_litellm_budget_table()
-
-            ### RESET MULTI-WINDOW BUDGETS ###
-            await self.reset_budget_windows()
+        await self.reset_budget_for_litellm_keys()
+        await self.reset_budget_for_litellm_users()
+        await self.reset_budget_for_litellm_teams()
+        await self.reset_budget_for_litellm_budget_table()
+        await self.reset_budget_windows()
 
     @staticmethod
     async def _invalidate_spend_counter(counter_key: str) -> None:
@@ -194,238 +291,195 @@ class ResetBudgetJob:
                 e,
             )
 
-    async def _cascade_reset_spend_for_budget_link(
+    async def _fetch_linked_rows(
         self,
-        budgets_to_reset: list[LiteLLM_BudgetTableFull],
         table: SpendLinkedTable[_RowT],
-        counter_key_fn: Callable[[_RowT], str],
+        where: Mapping[str, object],
         log_subject: str,
-        extra_where: dict[str, object] | None = None,
-        cache_key_fn: Callable[[_RowT], str | Sequence[str]] | None = None,
-    ):
-        """
-        Generic cascade: zero spend on rows whose budget_id is in the reset set.
+    ) -> tuple[_RowT, ...]:
+        """Read the rows the cascade will zero, so their counters can be
+        invalidated once the transaction commits."""
+        try:
+            return tuple(await table.find_many(where=where))
+        except Exception as e:
+            verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e)
+            return ()
 
-        ``cache_key_fn`` is optional: when provided, after the DB update each
-        matching row's entry or entries in ``user_api_key_cache`` are dropped so
-        cached spend cannot stay pinned above the zeroed DB row after a reset.
+    async def _collect_endusers_to_reset(self, budget_ids: Sequence[str]) -> tuple[_EndUserRow, ...]:
+        linked: Final[Sequence[_EndUserRow] | None] = await self.prisma_client.get_data(
+            table_name="enduser",
+            query_type="find_all",
+            budget_id_list=list(budget_ids),
+        )
+        if litellm.max_end_user_budget_id is None or litellm.max_end_user_budget_id not in budget_ids:
+            return tuple(linked or ())
+        return (*(linked or ()), *await self._get_endusers_with_no_budget_id())
+
+    async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade:
+        """Resolve every row the expiring budget tiers gate, before any write.
+
+        Keys carrying their own budget_duration are left out: they run on their
+        own schedule via reset_budget_for_litellm_keys(), so sweeping them here
+        would reset them twice.
         """
-        budget_ids: Final = [b.budget_id for b in budgets_to_reset if b.budget_id is not None]
+        budget_ids: Final = tuple(b.budget_id for b in budgets_to_reset if b.budget_id is not None)
         if not budget_ids:
+            return _EMPTY_CASCADE
+
+        team_memberships: Final[tuple[_TeamMembershipRow, ...]] = await self._fetch_linked_rows(
+            table=TeamMembershipRepository(self.prisma_client).table,
+            where=_budget_link_where(budget_ids),
+            log_subject="team memberships",
+        )
+        keys: Final[tuple[_KeyRow, ...]] = await self._fetch_linked_rows(
+            table=VerificationTokenRepository(self.prisma_client).table,
+            where=_budget_link_where(budget_ids, _LINKED_KEYS_WHERE),
+            log_subject="keys",
+        )
+        orgs: Final[tuple[_OrgRow, ...]] = await self._fetch_linked_rows(
+            table=OrganizationRepository(self.prisma_client).table,
+            where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
+            log_subject="orgs",
+        )
+        tags: Final[tuple[_TagRow, ...]] = await self._fetch_linked_rows(
+            table=TagRepository(self.prisma_client).table,
+            where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
+            log_subject="tags",
+        )
+        return _BudgetCascade(
+            budgets=tuple(budgets_to_reset),
+            budget_ids=budget_ids,
+            budget_resets=tuple(
+                (
+                    b.budget_id,
+                    compute_budget_reset_at(budget_duration=b.budget_duration, settings=self.reset_settings),
+                )
+                for b in budgets_to_reset
+                if b.budget_id is not None and b.budget_duration is not None
+            ),
+            endusers=await self._collect_endusers_to_reset(budget_ids),
+            counter_keys=(
+                *(_team_membership_counter_key(row) for row in team_memberships),
+                *(_key_counter_key(row) for row in keys),
+                *(_org_counter_key(row) for row in orgs),
+                *(_tag_counter_key(row) for row in tags),
+            ),
+            cache_keys=(
+                *(key for row in team_memberships for key in _team_membership_cache_keys(row)),
+                *(key for row in keys for key in _key_cache_keys(row)),
+                *(key for row in orgs for key in _org_cache_keys(row)),
+                *(key for row in tags for key in _tag_cache_keys(row)),
+            ),
+        )
+
+    async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None:
+        """Zero the gated spend and advance ``budget_reset_at`` in one transaction.
+
+        Advancing the window on its own hides the tier from every later tick
+        while its dependents stay pinned at the cap for the whole window;
+        batching both means a mid-cascade failure persists nothing and the rows
+        stay due for the next run.
+        """
+        if not cascade.budget_ids:
             return
 
-        where: Final[dict[str, object]] = {"budget_id": {"in": budget_ids}}
-        if extra_where:
-            where.update(extra_where)
+        enduser_ids: Final = tuple(row.user_id for row in cascade.endusers)
+        async with budget_cascade_unit_of_work(self.prisma_client.db.batch_) as uow:
+            uow.team_memberships.queue_spend_zero(where=_budget_link_where(cascade.budget_ids))
+            uow.keys.queue_spend_zero(where=_budget_link_where(cascade.budget_ids, _LINKED_KEYS_WHERE))
+            uow.organizations.queue_spend_zero(where=_budget_link_where(cascade.budget_ids, _SPENT_ROWS_WHERE))
+            uow.tags.queue_spend_zero(where=_budget_link_where(cascade.budget_ids, _SPENT_ROWS_WHERE))
+            if enduser_ids:
+                uow.endusers.queue_spend_zero(where={"user_id": {"in": list(enduser_ids)}})
+            for budget_id, budget_reset_at in cascade.budget_resets:
+                uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)
 
-        try:
-            rows: Sequence[_RowT] = await table.find_many(where=where)
-        except Exception as e:
-            rows = ()
-            verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e)
-
-        update_result: Final = await table.update_many(where=where, data={"spend": 0})
-
-        for row in rows:
-            await self._invalidate_spend_counter(counter_key_fn(row))
-            if cache_key_fn is not None:
-                cache_keys = cache_key_fn(row)
-                if isinstance(cache_keys, str):
-                    cache_keys = [cache_keys]
-                for cache_key in cache_keys:
-                    await self._invalidate_user_api_key_cache_entry(cache_key)
-
-        return update_result
-
-    async def reset_budget_for_litellm_team_members(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]):
-        """
-        Resets the budget for all LiteLLM Team Members if their budget has expired
-        """
-        return await self._cascade_reset_spend_for_budget_link(
-            budgets_to_reset=budgets_to_reset,
-            table=TeamMembershipRepository(self.prisma_client).table,
-            counter_key_fn=_team_membership_counter_key,
-            log_subject="team memberships",
-            cache_key_fn=_team_membership_cache_key,
-        )
-
-    async def reset_budget_for_keys_linked_to_budgets(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]):
-        """
-        Resets the spend for keys linked to budget tiers that are being reset.
-
-        Excludes keys with their own budget_duration; those are reset by
-        reset_budget_for_litellm_keys() to avoid double-resetting.
-        """
-        return await self._cascade_reset_spend_for_budget_link(
-            budgets_to_reset=budgets_to_reset,
-            table=VerificationTokenRepository(self.prisma_client).table,
-            counter_key_fn=_key_counter_key,
-            log_subject="keys",
-            extra_where={"budget_duration": None, "spend": {"gt": 0}},
-            cache_key_fn=_key_cache_key,
-        )
-
-    async def reset_budget_for_orgs_linked_to_budgets(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]):
-        """
-        Resets the spend for orgs linked to budget tiers that are being reset.
-        """
-        return await self._cascade_reset_spend_for_budget_link(
-            budgets_to_reset=budgets_to_reset,
-            table=OrganizationRepository(self.prisma_client).table,
-            counter_key_fn=_org_counter_key,
-            log_subject="orgs",
-            extra_where={"spend": {"gt": 0}},
-            cache_key_fn=_org_cache_keys,
-        )
-
-    async def reset_budget_for_tags_linked_to_budgets(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]):
-        """
-        Resets the spend for tags linked to budget tiers that are being reset.
-
-        Also drops each tag's ``user_api_key_cache`` entry so the next
-        ``_tag_max_budget_check`` reloads the zeroed row from the DB.
-        ``SpendCounterReseed.from_db`` intentionally returns ``None`` for
-        tags, so the budget check falls back to the cached
-        ``LiteLLM_TagTable.spend`` once the spend counter expires; without
-        this invalidation, that stale ``.spend`` keeps the tag over-budget
-        indefinitely.
-        """
-        return await self._cascade_reset_spend_for_budget_link(
-            budgets_to_reset=budgets_to_reset,
-            table=TagRepository(self.prisma_client).table,
-            counter_key_fn=_tag_counter_key,
-            log_subject="tags",
-            extra_where={"spend": {"gt": 0}},
-            cache_key_fn=_tag_cache_key,
-        )
-
-    async def reset_budget_for_litellm_budget_table(self):
-        """
-        Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
-        The corresponding Budget duration is also updated.
-        """
+    async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None:
+        for counter_key in cascade.counter_keys:
+            await self._invalidate_spend_counter(counter_key)
+        for cache_key in cascade.cache_keys:
+            await self._invalidate_user_api_key_cache_entry(cache_key)
 
+    async def _reset_expired_budget_cascade(self) -> _BudgetCascadeCommitted | _BudgetCascadeFailed:
         now: Final = datetime.now(timezone.utc)
-        start_time: Final = time.time()
-        endusers_to_reset: list[LiteLLM_EndUserTable] | None = None
-        budgets_to_reset: list[LiteLLM_BudgetTableFull] | None = None
-        updated_endusers: Final[list[LiteLLM_EndUserTable]] = []
-        failed_endusers: Final = []
         try:
-            budgets_to_reset = await self.prisma_client.get_data(
-                table_name="budget", query_type="find_all", reset_at=now
-            )
-
-            if budgets_to_reset is not None and len(budgets_to_reset) > 0:
-                for budget in budgets_to_reset:
-                    budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now, self.reset_settings)
-
-                await self.prisma_client.update_data(
-                    query_type="update_many",
-                    data_list=budgets_to_reset,
-                    table_name="budget",
-                )
-
-                budget_ids_to_reset = [budget.budget_id for budget in budgets_to_reset if budget.budget_id is not None]
-
-                endusers_to_reset = await self.prisma_client.get_data(
-                    table_name="enduser",
-                    query_type="find_all",
-                    budget_id_list=budget_ids_to_reset,
-                )
-
-                # Also reset end users with no budget_id (NULL) who use the
-                # default budget via litellm.max_end_user_budget_id.  These
-                # users are enforced in-memory but never had budget_id
-                # persisted, so the query above misses them.
-                if litellm.max_end_user_budget_id is not None and litellm.max_end_user_budget_id in budget_ids_to_reset:
-                    default_budget_endusers: Final = await self._get_endusers_with_no_budget_id()
-                    if default_budget_endusers:
-                        if endusers_to_reset is None:
-                            endusers_to_reset = default_budget_endusers
-                        else:
-                            endusers_to_reset.extend(default_budget_endusers)
-
-                await self.reset_budget_for_litellm_team_members(budgets_to_reset=budgets_to_reset)
-
-                await self.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)
-
-                await self.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=budgets_to_reset)
-
-                await self.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=budgets_to_reset)
-
-            if endusers_to_reset is not None and len(endusers_to_reset) > 0:
-                for enduser in endusers_to_reset:
-                    try:
-                        updated_enduser = await ResetBudgetJob._reset_budget_for_enduser(enduser=enduser)
-                        if updated_enduser is not None:
-                            updated_endusers.append(updated_enduser)
-                        else:
-                            failed_endusers.append(
-                                {
-                                    "enduser": enduser,
-                                    "error": "Returned None without exception",
-                                }
-                            )
-                    except Exception as e:
-                        failed_endusers.append({"enduser": enduser, "error": str(e)})
-                        verbose_proxy_logger.exception("Failed to reset budget for enduser: %s", enduser)
-
-                verbose_proxy_logger.debug(
-                    "Updated users %s",
-                    json.dumps(updated_endusers, indent=4, default=str),
-                )
-
-                await self.prisma_client.update_data(
-                    query_type="update_many",
-                    data_list=updated_endusers,
-                    table_name="enduser",
-                )
-
-            end_time = time.time()
-            if len(failed_endusers) > 0:  # If any endusers failed to reset
-                raise Exception(
-                    f"Failed to reset {len(failed_endusers)} endusers: {json.dumps(failed_endusers, default=str)}"
-                )
-
-            asyncio.create_task(
-                self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
-                    service=ServiceTypes.RESET_BUDGET_JOB,
-                    duration=end_time - start_time,
-                    call_type="reset_budget_budget_table",
-                    start_time=start_time,
-                    end_time=end_time,
-                    event_metadata={
-                        "num_budgets_found": (len(budgets_to_reset) if budgets_to_reset else 0),
-                        "budgets_found": json.dumps(budgets_to_reset, indent=4, default=str),
-                        "num_endusers_found": (len(endusers_to_reset) if endusers_to_reset else 0),
-                        "endusers_found": json.dumps(endusers_to_reset, indent=4, default=str),
-                        "num_endusers_updated": len(updated_endusers),
-                        "endusers_updated": json.dumps(updated_endusers, indent=4, default=str),
-                        "num_endusers_failed": len(failed_endusers),
-                        "endusers_failed": json.dumps(failed_endusers, indent=4, default=str),
-                    },
-                )
+            budgets_to_reset: Final[Sequence[LiteLLM_BudgetTableFull] | None] = await self.prisma_client.get_data(
+                table_name="budget",
+                query_type="find_all",
+                reset_at=now,
+                limit=RESET_BUDGET_JOB_BATCH_SIZE,
             )
+            cascade: Final = await self._collect_budget_cascade(budgets_to_reset or ())
         except Exception as e:
-            end_time = time.time()
-            asyncio.create_task(
-                self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
-                    service=ServiceTypes.RESET_BUDGET_JOB,
-                    duration=end_time - start_time,
-                    error=e,
-                    call_type="reset_budget_endusers",
-                    start_time=start_time,
-                    end_time=end_time,
-                    event_metadata={
-                        "num_budgets_found": (len(budgets_to_reset) if budgets_to_reset else 0),
-                        "budgets_found": json.dumps(budgets_to_reset, indent=4, default=str),
-                        "num_endusers_found": (len(endusers_to_reset) if endusers_to_reset else 0),
-                        "endusers_found": json.dumps(endusers_to_reset, indent=4, default=str),
-                    },
+            return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e)
+
+        try:
+            await self._commit_budget_cascade(cascade)
+        except Exception as e:
+            return _BudgetCascadeFailed(cascade=cascade, error=e)
+
+        await self._invalidate_budget_cascade_caches(cascade)
+        return _BudgetCascadeCommitted(
+            cascade=cascade,
+            advanced=_count_advanced(
+                (reset_at for _, reset_at in cascade.budget_resets),
+                cutoff=datetime.now(timezone.utc),
+            ),
+        )
+
+    async def reset_budget_for_litellm_budget_table(self) -> None:
+        """
+        Resets the spend a budget tier gates (end users, team members, keys,
+        orgs, tags) and advances the tier's budget_reset_at, atomically.
+
+        Caches are invalidated only after the transaction commits, so a failed
+        run cannot leave a zeroed counter in front of an un-reset DB row.
+        """
+        await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk)
+
+    async def _reset_budget_for_litellm_budget_table_chunk(self) -> _ChunkOutcome:
+        start_time: Final = time.time()
+        outcome: Final = await self._reset_expired_budget_cascade()
+        end_time: Final = time.time()
+
+        match outcome:
+            case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced):
+                asyncio.create_task(
+                    self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
+                        service=ServiceTypes.RESET_BUDGET_JOB,
+                        duration=end_time - start_time,
+                        call_type="reset_budget_budget_table",
+                        start_time=start_time,
+                        end_time=end_time,
+                        event_metadata={
+                            **_budget_cascade_event_metadata(cascade),
+                            "num_endusers_updated": len(cascade.endusers),
+                            "num_endusers_failed": 0,
+                        },
+                    )
                 )
-            )
-            verbose_proxy_logger.exception("Failed to reset budget for endusers: %s", e)
+                return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced)
+            case _BudgetCascadeFailed(cascade=cascade, error=error):
+                verbose_proxy_logger.exception(
+                    "Failed to reset the budget table cascade (team member, enduser, org and tag spend, plus "
+                    "budget_reset_at); nothing was committed and the budgets stay due for the next run: %s",
+                    error,
+                    exc_info=error,
+                )
+                asyncio.create_task(
+                    self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
+                        service=ServiceTypes.RESET_BUDGET_JOB,
+                        duration=end_time - start_time,
+                        error=error,
+                        call_type="reset_budget_endusers",
+                        start_time=start_time,
+                        end_time=end_time,
+                        event_metadata=_budget_cascade_event_metadata(cascade),
+                    )
+                )
+                return _NO_PROGRESS
+            case _:
+                assert_never(outcome)
 
     async def _get_endusers_with_no_budget_id(
         self,
@@ -486,18 +540,50 @@ class ResetBudgetJob:
             for t in updated_teams:
                 uow.teams.queue_spend_reset(team_id=t.team_id, budget_reset_at=t.budget_reset_at)
 
-    async def reset_budget_for_litellm_keys(self):
+    def _emit_phase_failure(
+        self,
+        call_type: str,
+        error: Exception,
+        start_time: float,
+        end_time: float,
+        event_metadata: dict[str, object],
+    ) -> None:
+        """Report rows that could not be reset without failing the chunk: the
+        rows that did reset are already committed, and raising here would cost
+        the phase every remaining chunk this tick.
+        """
+        verbose_proxy_logger.error("%s: %s", call_type, error)
+        asyncio.create_task(
+            self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
+                service=ServiceTypes.RESET_BUDGET_JOB,
+                duration=end_time - start_time,
+                error=error,
+                call_type=call_type,
+                start_time=start_time,
+                end_time=end_time,
+                event_metadata=event_metadata,
+            )
+        )
+
+    async def reset_budget_for_litellm_keys(self) -> None:
         """
         Resets the budget for all the litellm keys
 
         Catches Exceptions and logs them
         """
+        await _run_phase_in_chunks(self._reset_budget_for_litellm_keys_chunk)
+
+    async def _reset_budget_for_litellm_keys_chunk(self) -> _ChunkOutcome:
         now: Final = datetime.utcnow()
         start_time: Final = time.time()
         keys_to_reset: list[LiteLLM_VerificationToken] | None = None
         try:
             keys_to_reset = await self.prisma_client.get_data(
-                table_name="key", query_type="find_all", expires=now, reset_at=now
+                table_name="key",
+                query_type="find_all",
+                expires=now,
+                reset_at=now,
+                limit=RESET_BUDGET_JOB_BATCH_SIZE,
             )
             verbose_proxy_logger.debug("Keys to reset %s", json.dumps(keys_to_reset, indent=4, default=str))
             updated_keys: Final[list[LiteLLM_VerificationToken]] = []
@@ -528,8 +614,25 @@ class ResetBudgetJob:
                             await self._invalidate_spend_counter(f"spend:key:{token}")
 
             end_time = time.time()
-            if len(failed_keys) > 0:  # If any keys failed to reset
-                raise Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}")
+            outcome: Final = _ChunkOutcome(
+                fetched=len(keys_to_reset) if keys_to_reset else 0,
+                advanced=_count_advanced(
+                    (k.budget_reset_at for k in updated_keys),
+                    cutoff=datetime.now(timezone.utc),
+                ),
+            )
+            if len(failed_keys) > 0:
+                self._emit_phase_failure(
+                    call_type="reset_budget_keys",
+                    error=Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}"),
+                    start_time=start_time,
+                    end_time=end_time,
+                    event_metadata={
+                        "num_keys_found": len(keys_to_reset) if keys_to_reset else 0,
+                        "keys_found": json.dumps(keys_to_reset, indent=4, default=str),
+                    },
+                )
+                return outcome
 
             asyncio.create_task(
                 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
@@ -565,16 +668,27 @@ class ResetBudgetJob:
                 )
             )
             verbose_proxy_logger.exception("Failed to reset budget for keys: %s", e)
+            return _NO_PROGRESS
+        else:
+            return outcome
 
-    async def reset_budget_for_litellm_users(self):
+    async def reset_budget_for_litellm_users(self) -> None:
         """
         Resets the budget for all LiteLLM Internal Users if their budget has expired
         """
+        await _run_phase_in_chunks(self._reset_budget_for_litellm_users_chunk)
+
+    async def _reset_budget_for_litellm_users_chunk(self) -> _ChunkOutcome:
         now: Final = datetime.utcnow()
         start_time: Final = time.time()
         users_to_reset: list[LiteLLM_UserTable] | None = None
         try:
-            users_to_reset = await self.prisma_client.get_data(table_name="user", query_type="find_all", reset_at=now)
+            users_to_reset = await self.prisma_client.get_data(
+                table_name="user",
+                query_type="find_all",
+                reset_at=now,
+                limit=RESET_BUDGET_JOB_BATCH_SIZE,
+            )
             updated_users: Final[list[LiteLLM_UserTable]] = []
             failed_users: Final = []
             if users_to_reset is not None and len(users_to_reset) > 0:
@@ -609,8 +723,27 @@ class ResetBudgetJob:
                             await self._invalidate_global_proxy_spend_cache()
 
             end_time = time.time()
-            if len(failed_users) > 0:  # If any users failed to reset
-                raise Exception(f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}")
+            outcome: Final = _ChunkOutcome(
+                fetched=len(users_to_reset) if users_to_reset else 0,
+                advanced=_count_advanced(
+                    (u.budget_reset_at for u in updated_users),
+                    cutoff=datetime.now(timezone.utc),
+                ),
+            )
+            if len(failed_users) > 0:
+                self._emit_phase_failure(
+                    call_type="reset_budget_users",
+                    error=Exception(
+                        f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}"
+                    ),
+                    start_time=start_time,
+                    end_time=end_time,
+                    event_metadata={
+                        "num_users_found": len(users_to_reset) if users_to_reset else 0,
+                        "users_found": json.dumps(users_to_reset, indent=4, default=str),
+                    },
+                )
+                return outcome
 
             asyncio.create_task(
                 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
@@ -646,16 +779,27 @@ class ResetBudgetJob:
                 )
             )
             verbose_proxy_logger.exception("Failed to reset budget for users: %s", e)
+            return _NO_PROGRESS
+        else:
+            return outcome
 
-    async def reset_budget_for_litellm_teams(self):
+    async def reset_budget_for_litellm_teams(self) -> None:
         """
         Resets the budget for all LiteLLM Internal Teams if their budget has expired
         """
+        await _run_phase_in_chunks(self._reset_budget_for_litellm_teams_chunk)
+
+    async def _reset_budget_for_litellm_teams_chunk(self) -> _ChunkOutcome:
         now: Final = datetime.utcnow()
         start_time: Final = time.time()
         teams_to_reset: list[LiteLLM_TeamTable] | None = None
         try:
-            teams_to_reset = await self.prisma_client.get_data(table_name="team", query_type="find_all", reset_at=now)
+            teams_to_reset = await self.prisma_client.get_data(
+                table_name="team",
+                query_type="find_all",
+                reset_at=now,
+                limit=RESET_BUDGET_JOB_BATCH_SIZE,
+            )
             updated_teams: Final[list[LiteLLM_TeamTable]] = []
             failed_teams: Final = []
             if teams_to_reset is not None and len(teams_to_reset) > 0:
@@ -688,8 +832,27 @@ class ResetBudgetJob:
                             await self._invalidate_spend_counter(f"spend:team:{team_id}")
 
             end_time = time.time()
-            if len(failed_teams) > 0:  # If any teams failed to reset
-                raise Exception(f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}")
+            outcome: Final = _ChunkOutcome(
+                fetched=len(teams_to_reset) if teams_to_reset else 0,
+                advanced=_count_advanced(
+                    (t.budget_reset_at for t in updated_teams),
+                    cutoff=datetime.now(timezone.utc),
+                ),
+            )
+            if len(failed_teams) > 0:
+                self._emit_phase_failure(
+                    call_type="reset_budget_teams",
+                    error=Exception(
+                        f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}"
+                    ),
+                    start_time=start_time,
+                    end_time=end_time,
+                    event_metadata={
+                        "num_teams_found": len(teams_to_reset) if teams_to_reset else 0,
+                        "teams_found": json.dumps(teams_to_reset, indent=4, default=str),
+                    },
+                )
+                return outcome
 
             asyncio.create_task(
                 self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
@@ -725,6 +888,9 @@ class ResetBudgetJob:
                 )
             )
             verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e)
+            return _NO_PROGRESS
+        else:
+            return outcome
 
     @staticmethod
     async def _reset_expired_window(
@@ -882,33 +1048,6 @@ class ResetBudgetJob:
         )
         return user
 
-    @staticmethod
-    async def _reset_budget_for_enduser(
-        enduser: LiteLLM_EndUserTable,
-    ) -> LiteLLM_EndUserTable | None:
-        try:
-            enduser.spend = 0.0
-        except Exception as e:
-            verbose_proxy_logger.exception("Error resetting budget for enduser: %s. Item: %s", e, enduser)
-            raise e
-        return enduser
-
-    @staticmethod
-    async def _reset_budget_reset_at_date(
-        budget: LiteLLM_BudgetTableFull,
-        current_time: datetime,
-        reset_settings: BudgetResetSettings,
-    ) -> LiteLLM_BudgetTableFull:
-        try:
-            if budget.budget_duration is not None:
-                budget.budget_reset_at = compute_budget_reset_at(
-                    budget_duration=budget.budget_duration, settings=reset_settings
-                )
-        except Exception as e:
-            verbose_proxy_logger.exception("Error resetting budget_reset_at for budget: %s. Item: %s", e, budget)
-            raise e
-        return budget
-
     @staticmethod
     async def _reset_budget_for_key(
         key: LiteLLM_VerificationToken,
diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py
index 446ea76752e..8c6195388c5 100644
--- a/litellm/proxy/management_endpoints/budget_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py
@@ -20,7 +20,10 @@ from fastapi import APIRouter, Depends, HTTPException
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
-from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
+from litellm.proxy.management_endpoints.common_utils import (
+    _user_has_admin_view,
+    validate_budget_duration,
+)
 from litellm.proxy.utils import jsonify_object
 from litellm.repositories.budget_repository import BudgetRepository
 
@@ -72,6 +75,8 @@ async def new_budget(
             detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"},
         )
 
+    validate_budget_duration(budget_obj.budget_duration)
+
     # Validate model_max_budget if present
     if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0:
         from litellm.proxy.management_endpoints.key_management_endpoints import (
@@ -153,6 +158,8 @@ async def update_budget(
             detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"},
         )
 
+    validate_budget_duration(budget_obj.budget_duration)
+
     # Validate model_max_budget if present in update
     if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0:
         from litellm.proxy.management_endpoints.key_management_endpoints import (
diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py
index 3868b04f385..2241884faf1 100644
--- a/litellm/proxy/management_endpoints/common_utils.py
+++ b/litellm/proxy/management_endpoints/common_utils.py
@@ -22,6 +22,35 @@ def validate_finite_spend(spend: float | None) -> None:
         )
 
 
+def validate_budget_duration(budget_duration: str | None) -> None:
+    """Reject budget durations that can't be parsed, are non-positive, or
+    overflow date math, so a bad value can't be persisted and later crash the
+    budget reset job.
+
+    A non-positive duration also resolves to a reset time of "now", which leaves
+    the row permanently due: the reset job re-reads it every tick and, once
+    enough of them exist, they fill each batch and starve every other tenant's
+    reset.
+    """
+    if budget_duration is None:
+        return
+
+    from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+    from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+
+    try:
+        if duration_in_seconds(budget_duration) <= 0:
+            raise ValueError("budget_duration must be positive")
+        get_budget_reset_time(budget_duration=budget_duration)
+    except (ValueError, OverflowError):
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": f"Invalid budget_duration '{budget_duration}'. Use a format like '1h', '24h', '7d', or '30d'."
+            },
+        )
+
+
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import DualCache
 from litellm.proxy._types import (
diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py
index a51ff48aab6..bfc70da46ea 100644
--- a/litellm/proxy/management_endpoints/customer_endpoints.py
+++ b/litellm/proxy/management_endpoints/customer_endpoints.py
@@ -23,6 +23,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
+from litellm.proxy.management_endpoints.common_utils import validate_budget_duration
 from litellm.proxy.management_helpers.object_permission_utils import (
     _set_object_permission,
     handle_update_object_permission_common,
@@ -184,6 +185,7 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None:
 
     if budget_kv_pairs:
         budget_request: Final = BudgetNewRequest(**budget_kv_pairs)
+        validate_budget_duration(budget_request.budget_duration)
         if budget_request.budget_reset_at is None and budget_request.budget_duration is not None:
             budget_request.budget_reset_at = datetime.utcnow() + timedelta(
                 seconds=duration_in_seconds(duration=budget_request.budget_duration)
diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index abc5d3e53ff..a416a197ab8 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -42,6 +42,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     _is_user_team_admin,
     _user_has_admin_view,
     require_caller_user_id_for_non_admin,
+    validate_budget_duration,
     validate_finite_spend,
 )
 from litellm.proxy.management_endpoints.key_management_endpoints import (
@@ -506,6 +507,8 @@ async def new_user(
                 status_code=500,
                 detail=CommonProxyErrors.db_not_connected_error.value,
             )
+        validate_budget_duration(data.budget_duration)
+
         # Check for duplicate user_id or email
         await _check_duplicate_user_id(data.user_id, prisma_client)
         await _check_duplicate_user_email(data.user_email, prisma_client)
@@ -1185,6 +1188,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
     if "budget_duration" in non_default_values:
         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"]
         )
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 3a97558fcbe..38b5d755535 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -82,6 +82,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     _set_object_metadata_field,
     _team_member_has_permission,
     _user_has_admin_view,
+    validate_budget_duration,
     validate_finite_spend,
 )
 from litellm.proxy.management_endpoints.model_management_endpoints import (
@@ -844,6 +845,8 @@ async def _common_key_generation_helper(
         premium_user=premium_user,
     )
 
+    validate_budget_duration(data.budget_duration)
+
     if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
         raise HTTPException(
             status_code=403,
@@ -1024,7 +1027,7 @@ async def _common_key_generation_helper(
 
     # Only set budget_duration on key when explicitly provided. Keys with budget_id
     # but no explicit budget_duration follow their linked budget tier's schedule;
-    # reset_budget_for_keys_linked_to_budgets() resets them when the tier resets.
+    # reset_budget_for_litellm_budget_table() resets them when the tier resets.
     # This avoids duplicating budget_duration on keys so tier updates apply automatically.
     if "budget_duration" in data_json:
         data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
@@ -2401,6 +2404,7 @@ async def _validate_update_key_data(
     """Validate permissions and constraints for key update."""
     # Reject NaN/±inf spend before it can reach the DB / spend counter.
     validate_finite_spend(data.spend)
+    validate_budget_duration(data.budget_duration)
 
     _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
 
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index f5d0d63e311..60d3d650d00 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -96,6 +96,7 @@ from litellm.proxy.management_endpoints.common_utils import (
     _update_metadata_fields,
     _upsert_budget_and_membership,
     _user_has_admin_view,
+    validate_budget_duration,
 )
 from litellm.proxy.management_endpoints.organization_endpoints import (
     add_member_to_organization,
@@ -1258,6 +1259,9 @@ async def new_team(
                 detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
             )
 
+        validate_budget_duration(data.budget_duration)
+        validate_budget_duration(data.team_member_budget_duration)
+
         if data.soft_budget is not None:
             if data.max_budget is not None:
                 # If max_budget is set, soft_budget must be strictly lower than max_budget
@@ -1947,6 +1951,9 @@ async def update_team(
                 detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
             )
 
+        validate_budget_duration(data.budget_duration)
+        validate_budget_duration(data.team_member_budget_duration)
+
         existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
 
         if existing_team_row is None:
@@ -2979,7 +2986,7 @@ async def team_member_add(
     except HTTPException as e:
         raise e
 
-    _validate_budget_duration(data.budget_duration)
+    validate_budget_duration(data.budget_duration)
 
     prisma_client = cast(PrismaClient, prisma_client)
 
@@ -3282,29 +3289,6 @@ def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> dict[str, objec
     }
 
 
-def _validate_budget_duration(budget_duration: str | None) -> None:
-    """Reject budget durations that can't be parsed, are non-positive, or
-    overflow date math, so a bad value can't be persisted and later crash the
-    budget reset job."""
-    if budget_duration is None:
-        return
-
-    from litellm.litellm_core_utils.duration_parser import duration_in_seconds
-    from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
-
-    try:
-        if duration_in_seconds(budget_duration) <= 0:
-            raise ValueError("budget_duration must be positive")
-        get_budget_reset_time(budget_duration=budget_duration)
-    except (ValueError, OverflowError):
-        raise HTTPException(
-            status_code=400,
-            detail={
-                "error": f"Invalid budget_duration '{budget_duration}'. Use a format like '1h', '24h', '7d', or '30d'."
-            },
-        )
-
-
 @router.post(
     "/team/member_update",
     tags=["team management"],
@@ -3342,7 +3326,7 @@ async def team_member_update(
             detail={"error": "Either user_id or user_email needs to be passed in"},
         )
 
-    _validate_budget_duration(data.budget_duration)
+    validate_budget_duration(data.budget_duration)
 
     _existing_team_row: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
 
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 5f22ca021ac..dd0c57aa911 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -3486,13 +3486,15 @@ class PrismaClient:
                                 r.expires = r.expires.isoformat()
                 elif query_type == "find_all" and expires is not None and reset_at is not None:
                     response = await VerificationTokenRepository(self).table.find_many(
+                        take=limit,
                         where={
                             "OR": [
                                 {"expires": None},
                                 {"expires": {"gt": expires}},
                             ],
                             "budget_reset_at": {"lt": reset_at},
-                        }
+                            "NOT": {"budget_duration": None},
+                        },
                     )
                     if response is not None and len(response) > 0:
                         for r in response:
@@ -3542,6 +3544,7 @@ class PrismaClient:
                     response = await UserRepository(self).table.find_many(where=key_val)
                 elif query_type == "find_all" and reset_at is not None:
                     response = await UserRepository(self).table.find_many(
+                        take=limit,
                         where={
                             # A user seeded from default_internal_user_params
                             # (or created via /user/new without an explicit
@@ -3552,16 +3555,12 @@ class PrismaClient:
                             # of the row, silently exceeding max_budget. Treat a
                             # NULL budget_reset_at with a non-NULL budget_duration
                             # as due, matching the budget-table query below.
+                            "NOT": {"budget_duration": None},
                             "OR": [
-                                {
-                                    "AND": [
-                                        {"budget_reset_at": None},
-                                        {"NOT": {"budget_duration": None}},
-                                    ]
-                                },
+                                {"budget_reset_at": None},
                                 {"budget_reset_at": {"lt": reset_at}},
                             ],
-                        }
+                        },
                     )
                 elif query_type == "find_all" and user_id_list is not None:
                     response = await UserRepository(self).table.find_many(where={"user_id": {"in": user_id_list}})
@@ -3617,17 +3616,14 @@ class PrismaClient:
             elif table_name == "budget" and reset_at is not None:
                 if query_type == "find_all":
                     response = await BudgetRepository(self).table.find_many(
+                        take=limit,
                         where={
+                            "NOT": {"budget_duration": None},
                             "OR": [
-                                {
-                                    "AND": [
-                                        {"budget_reset_at": None},
-                                        {"NOT": {"budget_duration": None}},
-                                    ]
-                                },
+                                {"budget_reset_at": None},
                                 {"budget_reset_at": {"lt": reset_at}},
-                            ]
-                        }
+                            ],
+                        },
                     )
                     return response
 
@@ -3645,20 +3641,17 @@ class PrismaClient:
                     )
                 elif query_type == "find_all" and reset_at is not None:
                     response = await TeamRepository(self).table.find_many(
+                        take=limit,
                         where={
                             # Same NULL budget_reset_at gap as the user query
                             # above: a team with a budget_duration but no
                             # initialized budget_reset_at would never be reset.
+                            "NOT": {"budget_duration": None},
                             "OR": [
-                                {
-                                    "AND": [
-                                        {"budget_reset_at": None},
-                                        {"NOT": {"budget_duration": None}},
-                                    ]
-                                },
+                                {"budget_reset_at": None},
                                 {"budget_reset_at": {"lt": reset_at}},
                             ],
-                        }
+                        },
                     )
                 elif query_type == "find_all" and user_id is not None:
                     response = await TeamRepository(self).table.find_many(
diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py
index 4f020480f9e..e2e7f1fac73 100644
--- a/litellm/repositories/__init__.py
+++ b/litellm/repositories/__init__.py
@@ -70,10 +70,14 @@ from litellm.repositories.table_repositories import (
 )
 from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.unit_of_work import (
+    BudgetCascadeUnitOfWork,
+    BudgetWindowWrites,
     KeySpendResetWrites,
+    LinkedSpendResetWrites,
     SpendResetUnitOfWork,
     TeamSpendResetWrites,
     UserSpendResetWrites,
+    budget_cascade_unit_of_work,
     spend_reset_unit_of_work,
 )
 from litellm.repositories.user_repository import UserRepository
@@ -88,7 +92,9 @@ __all__ = [
     "AgentsRepository",
     "AuditLogRepository",
     "BatchTable",
+    "BudgetCascadeUnitOfWork",
     "BudgetRepository",
+    "BudgetWindowWrites",
     "CacheConfigRepository",
     "ClaudeCodePluginRepository",
     "ConfigOverridesRepository",
@@ -107,6 +113,7 @@ __all__ = [
     "InvitationLinkRepository",
     "JWTKeyMappingRepository",
     "KeySpendResetWrites",
+    "LinkedSpendResetWrites",
     "MCPServerRepository",
     "MCPToolsetRepository",
     "MCPUserCredentialsRepository",
@@ -149,5 +156,6 @@ __all__ = [
     "WorkflowEventRepository",
     "WorkflowMessageRepository",
     "WorkflowRunRepository",
+    "budget_cascade_unit_of_work",
     "spend_reset_unit_of_work",
 ]
diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py
index 6aff196ff10..055c68163f9 100644
--- a/litellm/repositories/prisma_protocols.py
+++ b/litellm/repositories/prisma_protocols.py
@@ -29,6 +29,8 @@ class SpendLinkedTable(Protocol[RowT_co]):
 class BatchTable(Protocol):
     def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
 
+    def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
+
 
 class PrismaBatch(Protocol):
     @property
@@ -40,4 +42,19 @@ class PrismaBatch(Protocol):
     @property
     def litellm_teamtable(self) -> BatchTable: ...
 
+    @property
+    def litellm_budgettable(self) -> BatchTable: ...
+
+    @property
+    def litellm_teammembership(self) -> BatchTable: ...
+
+    @property
+    def litellm_organizationtable(self) -> BatchTable: ...
+
+    @property
+    def litellm_tagtable(self) -> BatchTable: ...
+
+    @property
+    def litellm_endusertable(self) -> BatchTable: ...
+
     async def commit(self) -> None: ...
diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py
index 682e69d11eb..e504baceb9f 100644
--- a/litellm/repositories/unit_of_work.py
+++ b/litellm/repositories/unit_of_work.py
@@ -1,17 +1,21 @@
 """
-Unit of work over a single Prisma batch.
+Units of work over a single Prisma batch.
 
-``spend_reset_unit_of_work`` opens one ``db.batch_()`` and binds a typed write
+Each context manager here opens one ``db.batch_()`` and binds a typed write
 repository per table to it, so every update queued through the yielded object
 lands in the same transaction. The batch commits when the block exits cleanly
 and is abandoned, writing nothing, when the block raises.
 
-Each write repository queues narrow ``{spend, budget_reset_at}`` updates
+``spend_reset_unit_of_work`` covers the per-row key/user/team resets;
+``budget_cascade_unit_of_work`` covers a budget tier's reset, where the
+dependent spend and the tier's next window have to move together.
+
+Each write repository queues narrow ``{spend}`` / ``{budget_reset_at}`` updates
 instead of full-model writes, which trip ``prisma.errors.DataError`` on rows
 carrying fields the update input type rejects (see #27730).
 """
 
-from collections.abc import AsyncGenerator, Callable
+from collections.abc import AsyncGenerator, Callable, Mapping
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
@@ -43,6 +47,24 @@ class TeamSpendResetWrites:
         self.table.update(where={"team_id": team_id}, data={"spend": 0, "budget_reset_at": budget_reset_at})
 
 
+@dataclass(frozen=True, slots=True)
+class LinkedSpendResetWrites:
+    table: BatchTable
+
+    def queue_spend_zero(self, where: Mapping[str, object]) -> None:
+        self.table.update_many(where=where, data={"spend": 0})
+
+
+@dataclass(frozen=True, slots=True)
+class BudgetWindowWrites:
+    table: BatchTable
+
+    def queue_window_advance(self, budget_id: str, budget_reset_at: datetime) -> None:
+        """``update_many`` so a tier deleted between the read and the commit is a
+        no-op row count instead of a P2025 that aborts the whole chunk."""
+        self.table.update_many(where={"budget_id": budget_id}, data={"budget_reset_at": budget_reset_at})
+
+
 @dataclass(frozen=True, slots=True)
 class SpendResetUnitOfWork:
     keys: KeySpendResetWrites
@@ -50,6 +72,23 @@ class SpendResetUnitOfWork:
     teams: TeamSpendResetWrites
 
 
+@dataclass(frozen=True, slots=True)
+class BudgetCascadeUnitOfWork:
+    """Every write a budget-tier reset performs, bound to one batch.
+
+    The dependent spend rows and the budget rows' ``budget_reset_at`` advance
+    must land together: advancing the window without zeroing the spend it
+    gates leaves the dependents pinned at their cap until the next window.
+    """
+
+    team_memberships: LinkedSpendResetWrites
+    keys: LinkedSpendResetWrites
+    organizations: LinkedSpendResetWrites
+    tags: LinkedSpendResetWrites
+    endusers: LinkedSpendResetWrites
+    budgets: BudgetWindowWrites
+
+
 @asynccontextmanager
 async def spend_reset_unit_of_work(new_batch: Callable[[], PrismaBatch]) -> AsyncGenerator[SpendResetUnitOfWork, None]:
     batch = new_batch()
@@ -59,3 +98,19 @@ async def spend_reset_unit_of_work(new_batch: Callable[[], PrismaBatch]) -> Asyn
         teams=TeamSpendResetWrites(table=batch.litellm_teamtable),
     )
     await batch.commit()
+
+
+@asynccontextmanager
+async def budget_cascade_unit_of_work(
+    new_batch: Callable[[], PrismaBatch],
+) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]:
+    batch = new_batch()
+    yield BudgetCascadeUnitOfWork(
+        team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership),
+        keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken),
+        organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable),
+        tags=LinkedSpendResetWrites(table=batch.litellm_tagtable),
+        endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable),
+        budgets=BudgetWindowWrites(table=batch.litellm_budgettable),
+    )
+    await batch.commit()
diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json
index 4d1e73aab2d..da0f608fdb5 100644
--- a/ruff-strict-budget.json
+++ b/ruff-strict-budget.json
@@ -9,10 +9,10 @@
     "limit": 832
   },
   "ANN201": {
-    "limit": 2031
+    "limit": 2023
   },
   "ANN202": {
-    "limit": 861
+    "limit": 860
   },
   "ANN204": {
     "limit": 713
@@ -237,13 +237,13 @@
     "limit": 1226
   },
   "TRY002": {
-    "limit": 528
+    "limit": 524
   },
   "TRY004": {
     "limit": 96
   },
   "TRY201": {
-    "limit": 407
+    "limit": 405
   },
   "TRY203": {
     "limit": 113
diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py
index 00d5380b2f4..b13b7342c25 100644
--- a/tests/litellm_utils_tests/test_proxy_budget_reset.py
+++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py
@@ -44,39 +44,87 @@ def _attrify(d: dict):
     return _AttrDict(d)
 
 
-def _wire_batcher_for_test(prisma_client):
+def _wire_batcher_for_test(prisma_client, fail_commit=False):
     """
     Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is
-    awaitable and whose per-table .update() calls get captured. The reset job
-    writes key/user/team resets via prisma.db.batch_().<table>.update — not via
-    prisma_client.update_data — so tests must let that batch path complete.
+    awaitable and whose per-table .update()/.update_many() calls get captured.
+    The reset job writes every reset through prisma.db.batch_() — key/user/team
+    rows one by one, and the budget tier's cascade as a single transaction — so
+    tests must let that batch path complete.
 
-    Returns the list that will accumulate {table, where, data} dicts from
-    each captured update call.
+    Only committed batches contribute to the returned list, mirroring prisma:
+    with fail_commit=True the transaction blows up and must persist nothing.
+
+    Returns the list that will accumulate {table, op, where, data} dicts from
+    each captured write.
     """
     batch_calls = []
 
     def make_batcher():
+        queued = []
+
         class _Table:
             def __init__(self, table_name):
                 self._table_name = table_name
 
             def update(self, where=None, data=None):
-                batch_calls.append(
-                    {"table": self._table_name, "where": where, "data": data}
+                queued.append(
+                    {
+                        "table": self._table_name,
+                        "op": "update",
+                        "where": where,
+                        "data": data,
+                    }
                 )
 
+            def update_many(self, where=None, data=None):
+                queued.append(
+                    {
+                        "table": self._table_name,
+                        "op": "update_many",
+                        "where": where,
+                        "data": data,
+                    }
+                )
+
+        async def commit():
+            if fail_commit:
+                raise RuntimeError("simulated Postgres failure committing the batch")
+            batch_calls.extend(queued)
+
         batcher = MagicMock()
         batcher.litellm_verificationtoken = _Table("key")
         batcher.litellm_usertable = _Table("user")
         batcher.litellm_teamtable = _Table("team")
-        batcher.commit = AsyncMock(return_value=None)
+        batcher.litellm_budgettable = _Table("budget")
+        batcher.litellm_teammembership = _Table("team_membership")
+        batcher.litellm_organizationtable = _Table("org")
+        batcher.litellm_tagtable = _Table("tag")
+        batcher.litellm_endusertable = _Table("enduser")
+        batcher.commit = commit
         return batcher
 
     prisma_client.db.batch_ = MagicMock(side_effect=make_batcher)
     return batch_calls
 
 
+def _wire_cascade_reads_for_test(prisma_client):
+    """
+    The budget tier's cascade reads the rows it is about to zero, so their
+    spend counters can be invalidated after the commit. Give each of those
+    tables an awaitable find_many so the reads resolve instead of falling into
+    the job's warn-and-continue path.
+    """
+    for table in (
+        "litellm_teammembership",
+        "litellm_verificationtoken",
+        "litellm_organizationtable",
+        "litellm_tagtable",
+        "litellm_endusertable",
+    ):
+        getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[])
+
+
 @pytest.mark.asyncio
 async def test_reset_budget_keys_partial_failure():
     """
@@ -250,41 +298,18 @@ async def test_reset_budget_users_partial_failure():
 
 
 @pytest.mark.asyncio
-async def test_reset_budget_endusers_partial_failure():
+async def test_reset_budget_endusers_cascade_failure_is_all_or_nothing():
     """
-    Test that if one enduser fails to reset, the reset loop still processes the other endusers.
-    We simulate six endsers where the first fails and the others are updated.
+    A failure anywhere in the budget-tier cascade must persist nothing, so the
+    tier stays due and the next scheduler tick retries it. Before the fix the
+    job committed the new budget_reset_at first and zeroed the dependent spend
+    afterwards, so a failure here left the tier stamped for the next window
+    while every end user stayed at the cap.
     """
-    user1 = {
-        "user_id": "user1",
-        "spend": 20.0,
-        "budget_id": "budget1",
-    }  # Will trigger simulated failure
-    user2 = {
-        "user_id": "user2",
-        "spend": 25.0,
-        "budget_id": "budget1",
-    }  # Should be updated
-    user3 = {
-        "user_id": "user3",
-        "spend": 30.0,
-        "budget_id": "budget1",
-    }  # Should be updated
-    user4 = {
-        "user_id": "user4",
-        "spend": 35.0,
-        "budget_id": "budget1",
-    }  # Should be updated
-    user5 = {
-        "user_id": "user5",
-        "spend": 40.0,
-        "budget_id": "budget1",
-    }  # Should be updated
-    user6 = {
-        "user_id": "user6",
-        "spend": 45.0,
-        "budget_id": "budget1",
-    }  # Should be updated
+    endusers = [
+        _attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"})
+        for i in range(1, 7)
+    ]
 
     budget1 = LiteLLM_BudgetTableFull(
         **{
@@ -301,23 +326,13 @@ async def test_reset_budget_endusers_partial_failure():
         if table_name == "budget":
             return [budget1]
         elif table_name == "enduser":
-            return [user1, user2, user3, user4, user5, user6]
+            return endusers
         return []
 
     prisma_client.get_data = AsyncMock()
     prisma_client.get_data.side_effect = get_data_mock
-
     prisma_client.update_data = AsyncMock()
-    # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets)
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
+    batch_calls = _wire_batcher_for_test(prisma_client, fail_commit=True)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -326,41 +341,13 @@ async def test_reset_budget_endusers_partial_failure():
 
     job = ResetBudgetJob(proxy_logging_obj, prisma_client)
 
-    async def fake_reset_enduser(enduser):
-        if enduser["user_id"] == "user1":
-            raise Exception("Simulated failure for user1")
-        enduser["spend"] = 0.0
-        return enduser
+    await job.reset_budget_for_litellm_budget_table()
+    await asyncio.sleep(0.1)
 
-    async def fake_reset_team_members(budgets_to_reset):
-        return 1
-
-    with (
-        patch.object(
-            ResetBudgetJob,
-            "_reset_budget_for_enduser",
-            side_effect=fake_reset_enduser,
-        ) as mock_reset_enduser,
-        patch.object(
-            ResetBudgetJob,
-            "reset_budget_for_litellm_team_members",
-            side_effect=fake_reset_team_members,
-        ) as mock_reset_team_members,
-    ):
-        await job.reset_budget_for_litellm_budget_table()
-        await asyncio.sleep(0.1)
-
-    assert mock_reset_enduser.call_count == 6
-    assert prisma_client.update_data.await_count == 2
-    update_call = prisma_client.update_data.call_args
-    assert update_call.kwargs.get("table_name") == "enduser"
-    updated_users = update_call.kwargs.get("data_list", [])
-    assert len(updated_users) == 5
-    assert updated_users[0]["user_id"] == "user2"
-    assert updated_users[1]["user_id"] == "user3"
-    assert updated_users[2]["user_id"] == "user4"
-    assert updated_users[3]["user_id"] == "user5"
-    assert updated_users[4]["user_id"] == "user6"
+    assert batch_calls == [], "a failed cascade must not persist any write"
+    assert (
+        prisma_client.update_data.await_count == 0
+    ), "budget_reset_at must not be advanced outside the cascade transaction"
 
     failure_hook_calls = (
         proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
@@ -369,6 +356,66 @@ async def test_reset_budget_endusers_partial_failure():
         call.kwargs.get("call_type") == "reset_budget_endusers"
         for call in failure_hook_calls
     )
+    proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance():
+    """
+    The happy path: every end user the tier gates is zeroed and the tier's
+    budget_reset_at advances, all inside one transaction.
+    """
+    endusers = [
+        _attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"})
+        for i in range(1, 7)
+    ]
+
+    budget1 = LiteLLM_BudgetTableFull(
+        **{
+            "budget_id": "budget1",
+            "max_budget": 65.0,
+            "budget_duration": "2d",
+            "created_at": datetime.now(timezone.utc) - timedelta(days=3),
+        }
+    )
+
+    prisma_client = MagicMock()
+
+    async def get_data_mock(table_name, *args, **kwargs):
+        if table_name == "budget":
+            return [budget1]
+        elif table_name == "enduser":
+            return endusers
+        return []
+
+    prisma_client.get_data = AsyncMock()
+    prisma_client.get_data.side_effect = get_data_mock
+    prisma_client.update_data = AsyncMock()
+    batch_calls = _wire_batcher_for_test(prisma_client)
+
+    proxy_logging_obj = MagicMock()
+    proxy_logging_obj.service_logging_obj = MagicMock()
+    proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
+    proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
+
+    job = ResetBudgetJob(proxy_logging_obj, prisma_client)
+
+    await job.reset_budget_for_litellm_budget_table()
+    await asyncio.sleep(0.1)
+
+    assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction"
+
+    enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
+    assert len(enduser_writes) == 1
+    assert enduser_writes[0]["where"]["user_id"]["in"] == [f"user{i}" for i in range(1, 7)]
+    assert enduser_writes[0]["data"] == {"spend": 0}
+
+    budget_writes = [c for c in batch_calls if c["table"] == "budget"]
+    assert len(budget_writes) == 1
+    assert budget_writes[0]["where"] == {"budget_id": "budget1"}
+    assert budget_writes[0]["data"]["budget_reset_at"] > datetime.now(timezone.utc)
+
+    proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -500,16 +547,8 @@ async def test_reset_budget_continues_other_categories_on_failure():
     key1, key2 = _attrify(key1), _attrify(key2)
     user1, user2 = _attrify(user1), _attrify(user2)
     team1, team2 = _attrify(team1), _attrify(team2)
-    # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets)
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
+    enduser1 = _attrify(enduser1)
+    _wire_cascade_reads_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -541,13 +580,6 @@ async def test_reset_budget_continues_other_categories_on_failure():
         ).isoformat()
         return team
 
-    async def fake_reset_enduser(enduser):
-        enduser["spend"] = 0.0
-        return enduser
-
-    async def fake_reset_team_members(budgets_to_reset):
-        return 1
-
     with (
         patch.object(
             ResetBudgetJob, "_reset_budget_for_key", side_effect=fake_reset_key
@@ -558,14 +590,6 @@ async def test_reset_budget_continues_other_categories_on_failure():
         patch.object(
             ResetBudgetJob, "_reset_budget_for_team", side_effect=fake_reset_team
         ) as mock_reset_team,
-        patch.object(
-            ResetBudgetJob, "_reset_budget_for_enduser", side_effect=fake_reset_enduser
-        ) as mock_reset_enduser,
-        patch.object(
-            ResetBudgetJob,
-            "reset_budget_for_litellm_team_members",
-            side_effect=fake_reset_team_members,
-        ) as mock_reset_team_members,
     ):
         # Call the overall reset_budget method.
         await job.reset_budget()
@@ -575,29 +599,22 @@ async def test_reset_budget_continues_other_categories_on_failure():
     called_tables = {
         call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list
     }
-    if mock_reset_team_members.call_count > 0:
-        called_tables.add("team_membership")
-    assert called_tables == {
-        "key",
-        "user",
-        "team",
-        "budget",
-        "enduser",
-        "team_membership",
-    }
+    assert called_tables == {"key", "user", "team", "budget", "enduser"}
 
-    # After the fix, keys/users/teams write via prisma.db.batch_().<table>.update,
-    # so only budget + enduser still go through update_data.
-    calls = prisma_client.update_data.await_args_list
-    update_data_tables = [c.kwargs.get("table_name") for c in calls]
-    assert sorted(update_data_tables) == ["budget", "enduser"]
+    # Every category writes through the batch path now, so update_data is unused.
+    prisma_client.update_data.assert_not_awaited()
 
-    # Check enduser update: enduser succeed.
-    enduser_call = next(c for c in calls if c.kwargs.get("table_name") == "enduser")
-    assert len(enduser_call.kwargs.get("data_list", [])) == 1
+    # The budget tier's cascade still ran despite the failing user category.
+    assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1
+    enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
+    assert len(enduser_writes) == 1
+    assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}}
+    assert enduser_writes[0]["data"] == {"spend": 0}
 
     # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams.
-    key_writes = [c for c in batch_calls if c["table"] == "key"]
+    # `op` separates the per-row resets from the cascade sweep, which also
+    # targets the key table.
+    key_writes = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"]
     user_writes = [c for c in batch_calls if c["table"] == "user"]
     team_writes = [c for c in batch_calls if c["table"] == "team"]
     assert len(key_writes) == 2
@@ -974,12 +991,12 @@ async def test_service_logger_teams_failure():
 @pytest.mark.asyncio
 async def test_service_logger_endusers_success():
     """
-    Test that when resetting endusers succeeds the service logger success hook is called with
-    the correct metadata and no exception is logged.
+    Test that when the budget-tier cascade commits, the service logger success
+    hook is called with the correct metadata and no exception is logged.
     """
     endusers = [
-        {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"},
-        {"user_id": "user2", "spend": 25.0, "budget_id": "budget1"},
+        _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}),
+        _attrify({"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}),
     ]
     budgets = [
         LiteLLM_BudgetTableFull(
@@ -1002,16 +1019,8 @@ async def test_service_logger_endusers_success():
     prisma_client = MagicMock()
     prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
     prisma_client.update_data = AsyncMock()
-    # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets)
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
+    batch_calls = _wire_batcher_for_test(prisma_client)
+    _wire_cascade_reads_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -1020,31 +1029,16 @@ async def test_service_logger_endusers_success():
 
     job = ResetBudgetJob(proxy_logging_obj, prisma_client)
 
-    async def fake_reset_enduser(enduser):
-        enduser["spend"] = 0.0
-        return enduser
+    with patch(
+        "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
+    ) as mock_verbose_exc:
+        await job.reset_budget_for_litellm_budget_table()
+        await asyncio.sleep(0.1)
+        mock_verbose_exc.assert_not_called()
 
-    async def fake_reset_team_members(budgets_to_reset):
-        return 1
-
-    with (
-        patch.object(
-            ResetBudgetJob,
-            "_reset_budget_for_enduser",
-            side_effect=fake_reset_enduser,
-        ) as mock_reset_enduser,
-        patch.object(
-            ResetBudgetJob,
-            "reset_budget_for_litellm_team_members",
-            side_effect=fake_reset_team_members,
-        ) as mock_reset_team_members,
-    ):
-        with patch(
-            "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
-        ) as mock_verbose_exc:
-            await job.reset_budget_for_litellm_budget_table()
-            await asyncio.sleep(0.1)
-            mock_verbose_exc.assert_not_called()
+    enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
+    assert len(enduser_writes) == 1
+    assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1", "user2"]}}
 
     proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once()
     (
@@ -1062,12 +1056,12 @@ async def test_service_logger_endusers_success():
 @pytest.mark.asyncio
 async def test_service_logger_endusers_failure():
     """
-    Test that a failure during enduser reset calls the failure hook with appropriate metadata,
-    logs the exception, and does not call the success hook.
+    Test that a failed cascade calls the failure hook with the rows it had
+    found, logs the exception, and does not call the success hook.
     """
     endusers = [
-        {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"},
-        {"user_id": "user2", "spend": 25.0, "budget_id": "budget1"},
+        _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}),
+        _attrify({"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}),
     ]
     budgets = [
         LiteLLM_BudgetTableFull(
@@ -1090,16 +1084,8 @@ async def test_service_logger_endusers_failure():
     prisma_client = MagicMock()
     prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
     prisma_client.update_data = AsyncMock()
-    # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets)
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
+    _wire_batcher_for_test(prisma_client, fail_commit=True)
+    _wire_cascade_reads_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -1108,39 +1094,16 @@ async def test_service_logger_endusers_failure():
 
     job = ResetBudgetJob(proxy_logging_obj, prisma_client)
 
-    async def fake_reset_enduser(enduser):
-        if enduser["user_id"] == "user1":
-            raise Exception("Simulated failure for user1")
-        enduser["spend"] = 0.0
-        return enduser
-
-    async def fake_reset_team_members(budgets_to_reset):
-        return 1
-
-    with (
-        patch.object(
-            ResetBudgetJob,
-            "_reset_budget_for_enduser",
-            side_effect=fake_reset_enduser,
-        ) as mock_reset_enduser,
-        patch.object(
-            ResetBudgetJob,
-            "reset_budget_for_litellm_team_members",
-            side_effect=fake_reset_team_members,
-        ) as mock_reset_team_members,
-    ):
-        with patch(
-            "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
-        ) as mock_verbose_exc:
-            await job.reset_budget_for_litellm_budget_table()
-            await asyncio.sleep(0.1)
-            # Verify exception logging
-            assert mock_verbose_exc.call_count >= 1
-            # Verify exception was logged with correct message
-            assert any(
-                "Failed to reset budget for enduser" in str(call.args)
-                for call in mock_verbose_exc.call_args_list
-            )
+    with patch(
+        "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
+    ) as mock_verbose_exc:
+        await job.reset_budget_for_litellm_budget_table()
+        await asyncio.sleep(0.1)
+        # The log must name the whole cascade, not just end users: the write
+        # that failed could have been any of team member / enduser / org / tag
+        # spend or the budget_reset_at advance.
+        assert mock_verbose_exc.call_count == 1
+        assert "budget table cascade" in str(mock_verbose_exc.call_args.args[0])
 
     proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_called_once()
     (
@@ -1158,8 +1121,8 @@ async def test_service_logger_endusers_failure():
 @pytest.mark.asyncio
 async def test_reset_budget_for_litellm_team_members_called():
     """
-    Test that when reset_budget_for_litellm_budget_table is called,
-    team members' budgets are also reset via reset_budget_for_litellm_team_members
+    Test that when reset_budget_for_litellm_budget_table is called, team
+    members' spend is zeroed as part of the cascade transaction.
     """
     # Arrange
     budget1 = LiteLLM_BudgetTableFull(
@@ -1171,7 +1134,7 @@ async def test_reset_budget_for_litellm_team_members_called():
         }
     )
 
-    enduser1 = {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}
+    enduser1 = _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"})
 
     prisma_client = MagicMock()
 
@@ -1184,20 +1147,9 @@ async def test_reset_budget_for_litellm_team_members_called():
 
     prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
     prisma_client.update_data = AsyncMock()
-
-    # Mock the db.litellm_teammembership.update_many call
     prisma_client.db = MagicMock()
-    prisma_client.db.litellm_teammembership = MagicMock()
-    prisma_client.db.litellm_teammembership.update_many = AsyncMock(
-        return_value={"count": 2}
-    )
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
-        return_value={"count": 0}
-    )
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
+    batch_calls = _wire_batcher_for_test(prisma_client)
+    _wire_cascade_reads_for_test(prisma_client)
 
     proxy_logging_obj = MagicMock()
     proxy_logging_obj.service_logging_obj = MagicMock()
@@ -1206,23 +1158,11 @@ async def test_reset_budget_for_litellm_team_members_called():
 
     job = ResetBudgetJob(proxy_logging_obj, prisma_client)
 
-    async def fake_reset_enduser(enduser):
-        enduser["spend"] = 0.0
-        return enduser
-
-    with patch.object(
-        ResetBudgetJob,
-        "_reset_budget_for_enduser",
-        side_effect=fake_reset_enduser,
-    ):
-        # Act
-        await job.reset_budget_for_litellm_budget_table()
+    # Act
+    await job.reset_budget_for_litellm_budget_table()
 
     # Assert
-    # Verify that the team membership update was called
-    prisma_client.db.litellm_teammembership.update_many.assert_called_once()
-
-    # Verify the call was made with correct parameters
-    call_args = prisma_client.db.litellm_teammembership.update_many.call_args
-    assert call_args.kwargs["where"]["budget_id"]["in"] == ["budget1"]
-    assert call_args.kwargs["data"]["spend"] == 0
+    team_member_writes = [c for c in batch_calls if c["table"] == "team_membership"]
+    assert len(team_member_writes) == 1
+    assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"]
+    assert team_member_writes[0]["data"] == {"spend": 0}
diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
index 616ad8a0981..608dc8cb5c8 100644
--- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
+++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
@@ -2,7 +2,6 @@ import asyncio
 import json
 import os
 import sys
-import time
 import types
 from datetime import datetime, timedelta, timezone
 from datetime import time as dt_time
@@ -13,33 +12,19 @@ import pytest
 
 sys.path.insert(0, os.path.abspath("../../.."))  # Adds the parent directory to the system path
 
-from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import LiteLLM_VerificationToken
+from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module
 from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
 from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings
-from litellm.proxy.utils import ProxyLogging
 
 
 # Mock classes for testing
-class MockLiteLLMTeamMembership:
-    async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]:
-        # Mock the update_many method for litellm_teammembership
-        return {"count": 1}
+class MockTable:
+    """A single prisma table: records reads/writes and replays canned rows."""
 
-
-class MockLiteLLMVerificationToken:
     def __init__(self):
-        self.update_many_calls: List[Dict[str, Any]] = []
-
-    async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]:
-        self.update_many_calls.append({"where": where, "data": data})
-        return {"count": 1}
-
-
-class MockLiteLLMOrganizationTable:
-    def __init__(self):
-        self.update_many_calls: List[Dict[str, Any]] = []
         self.find_many_calls: List[Dict[str, Any]] = []
+        self.update_many_calls: List[Dict[str, Any]] = []
         self._find_many_results: List[Any] = []
 
     def set_find_many_results(self, results: List[Any]):
@@ -54,43 +39,12 @@ class MockLiteLLMOrganizationTable:
         return {"count": 1}
 
 
-class MockLiteLLMTagTable:
-    def __init__(self):
-        self.update_many_calls: List[Dict[str, Any]] = []
-        self.find_many_calls: List[Dict[str, Any]] = []
-        self._find_many_results: List[Any] = []
-
-    def set_find_many_results(self, results: List[Any]):
-        self._find_many_results = results
-
-    async def find_many(self, where: Dict[str, Any]) -> List[Any]:
-        self.find_many_calls.append({"where": where})
-        return self._find_many_results
-
-    async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]:
-        self.update_many_calls.append({"where": where, "data": data})
-        return {"count": 1}
-
-
-class MockLiteLLMEndUserTable:
-    def __init__(self):
-        self.find_many_calls: List[Dict[str, Any]] = []
-        self._find_many_results: List[Any] = []
-
-    def set_find_many_results(self, results: List[Any]):
-        self._find_many_results = results
-
-    async def find_many(self, where: Dict[str, Any]) -> List[Any]:
-        self.find_many_calls.append({"where": where})
-        return self._find_many_results
-
-
 class MockBatcher:
-    """Captures per-row update calls and exposes them after commit().
+    """Captures the writes queued on one `db.batch_()` and whether it committed.
 
-    Mirrors prisma's `db.batch_()` ergonomics enough that the reset job's
-    narrow-write helpers (`_write_key_reset_updates` et al) can run against
-    the mock and the test can assert on what would have been written.
+    Mirrors prisma's batch ergonomics enough that the reset job's write helpers
+    can run against the mock, and keeps `committed` so tests can prove a failed
+    cascade persisted nothing.
     """
 
     def __init__(self):
@@ -102,12 +56,23 @@ class MockBatcher:
                 _self._table_name = table_name
                 _self._outer = outer
 
+            def _record(_self, op, where, data):
+                _self._outer.calls.append({"table": _self._table_name, "op": op, "where": where, "data": data})
+
             def update(_self, where, data):
-                _self._outer.calls.append({"table": _self._table_name, "where": where, "data": data})
+                _self._record("update", where, data)
+
+            def update_many(_self, where, data):
+                _self._record("update_many", where, data)
 
         self.litellm_verificationtoken = _Table("key", self)
         self.litellm_usertable = _Table("user", self)
         self.litellm_teamtable = _Table("team", self)
+        self.litellm_budgettable = _Table("budget", self)
+        self.litellm_teammembership = _Table("team_membership", self)
+        self.litellm_organizationtable = _Table("org", self)
+        self.litellm_tagtable = _Table("tag", self)
+        self.litellm_endusertable = _Table("enduser", self)
 
     async def commit(self):
         self.committed = True
@@ -116,16 +81,20 @@ class MockBatcher:
 
 class MockDB:
     def __init__(self):
-        self.litellm_teammembership = MockLiteLLMTeamMembership()
-        self.litellm_verificationtoken = MockLiteLLMVerificationToken()
-        self.litellm_endusertable = MockLiteLLMEndUserTable()
-        self.litellm_organizationtable = MockLiteLLMOrganizationTable()
-        self.litellm_tagtable = MockLiteLLMTagTable()
+        self.litellm_teammembership = MockTable()
+        self.litellm_verificationtoken = MockTable()
+        self.litellm_endusertable = MockTable()
+        self.litellm_organizationtable = MockTable()
+        self.litellm_tagtable = MockTable()
         self.batch_calls: List[Dict[str, Any]] = []
+        self.batchers: List[MockBatcher] = []
 
     def batch_(self):
         batcher = MockBatcher()
-        # Aggregate calls across all batches so tests can assert on cumulative writes.
+        self.batchers.append(batcher)
+        # Aggregate calls across all batches so tests can assert on cumulative
+        # writes. Only committed batches contribute: an abandoned batch writes
+        # nothing, exactly as prisma behaves.
         original_commit = batcher.commit
 
         async def _record_and_commit():
@@ -152,9 +121,11 @@ class MockPrismaClient:
             "budget": [],
             "enduser": [],
         }
+        self.get_data_calls: List[Dict[str, Any]] = []
         self.db = MockDB()
 
     async def get_data(self, table_name, query_type, **kwargs):
+        self.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs})
         data = self.data.get(table_name, [])
 
         # Handle specific filtering for budget table queries
@@ -218,6 +189,39 @@ async def run_async_test(coro):
     return await coro
 
 
+_ALREADY_EXPIRED = object()
+
+
+def _budget_row(
+    budget_id: str = "test-budget-1",
+    budget_duration: Any = "7d",
+    budget_reset_at: Any = _ALREADY_EXPIRED,
+    max_budget: float = 10.0,
+):
+    """An expiring budget tier, shaped like the rows get_data() hands back."""
+    now = datetime.now(timezone.utc)
+    return type(
+        "LiteLLM_BudgetTableFull",
+        (),
+        {
+            "max_budget": max_budget,
+            "budget_duration": budget_duration,
+            "budget_reset_at": (now - timedelta(hours=1) if budget_reset_at is _ALREADY_EXPIRED else budget_reset_at),
+            "budget_id": budget_id,
+            "created_at": now - timedelta(days=30),
+        },
+    )
+
+
+def _batch_writes(mock_prisma_client, table: str, op: str | None = None) -> List[Dict[str, Any]]:
+    """Writes that were committed to the DB, optionally narrowed to one op."""
+    return [
+        call
+        for call in mock_prisma_client.db.batch_calls
+        if call["table"] == table and (op is None or call["op"] == op)
+    ]
+
+
 # Tests
 def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(reset_budget_job, mock_prisma_client):
     """A key with token=None must be skipped, not queued as where={"token": None}.
@@ -234,10 +238,10 @@ def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(rese
 
     asyncio.run(reset_budget_job._write_key_reset_updates(updated_keys=keys))
 
-    key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"]
-    assert key_writes == [
+    assert _batch_writes(mock_prisma_client, "key") == [
         {
             "table": "key",
+            "op": "update",
             "where": {"token": "tok-ok"},
             "data": {"spend": 0, "budget_reset_at": reset_at},
         }
@@ -369,18 +373,9 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
 
 
 def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client):
-    # Setup test data
+    """End-user spend is zeroed and the tier's window advances, in one batch."""
     now = datetime.now(timezone.utc)
-    test_budget = type(
-        "LiteLLM_BudgetTable",
-        (),
-        {
-            "max_budget": 500.0,
-            "budget_duration": "1d",
-            "budget_reset_at": now,
-            "budget_id": "test-budget-1",
-        },
-    )
+    test_budget = _budget_row(budget_id="test-budget-1", budget_duration="1d", budget_reset_at=now)
 
     test_enduser = type(
         "LiteLLM_EndUserTable",
@@ -395,16 +390,22 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client):
     mock_prisma_client.data["budget"] = [test_budget]
     mock_prisma_client.data["enduser"] = [test_enduser]
 
-    # Run the test
     asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    # Verify results
-    assert len(mock_prisma_client.updated_data["enduser"]) == 1
-    assert len(mock_prisma_client.updated_data["budget"]) == 1
-    updated_enduser = mock_prisma_client.updated_data["enduser"][0]
-    updated_budget = mock_prisma_client.updated_data["budget"][0]
-    assert updated_enduser.spend == 0.0
-    assert updated_budget.budget_reset_at > now
+    assert _batch_writes(mock_prisma_client, "enduser") == [
+        {
+            "table": "enduser",
+            "op": "update_many",
+            "where": {"user_id": {"in": ["test-enduser-1"]}},
+            "data": {"spend": 0},
+        }
+    ]
+
+    budget_writes = _batch_writes(mock_prisma_client, "budget")
+    assert len(budget_writes) == 1
+    assert budget_writes[0]["where"] == {"budget_id": "test-budget-1"}
+    assert budget_writes[0]["data"]["budget_reset_at"] > now
+    assert set(budget_writes[0]["data"].keys()) == {"budget_reset_at"}
 
 
 def test_reset_budget_all(reset_budget_job, mock_prisma_client):
@@ -485,190 +486,81 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
         ("user", {"user_id": "uid-all-1"}),
         ("team", {"team_id": "tid-all-1"}),
     ]:
-        writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == table_name]
+        writes = _batch_writes(mock_prisma_client, table_name, op="update")
         assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}"
         assert writes[0]["where"] == where
         assert writes[0]["data"]["spend"] == 0
         assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
 
-    # Enduser + budget rows still go through update_data (not narrowed; different path).
-    assert len(mock_prisma_client.updated_data["enduser"]) == 1
-    assert len(mock_prisma_client.updated_data["budget"]) == 1
-    assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0
-
-
-def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_client):
-    """
-    Test that when a budget tier is reset, keys linked to that budget
-    (via budget_id) that don't have their own budget_duration also get
-    their spend reset.
-
-    This covers the case where keys were created with budget_id but
-    budget_duration was not inherited to the key (pre-fix keys).
-    """
-    from litellm.proxy._types import LiteLLM_BudgetTableFull
-
-    now = datetime.now(timezone.utc)
-
-    # Create a budget tier that is due for reset
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
+    # The budget tier's cascade rides the same batch machinery.
+    assert _batch_writes(mock_prisma_client, "enduser") == [
         {
-            "max_budget": 10.0,
-            "budget_duration": "7d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "7d-budget-tier",
-            "created_at": now - timedelta(days=7),
+            "table": "enduser",
+            "op": "update_many",
+            "where": {"user_id": {"in": ["test-enduser-1"]}},
+            "data": {"spend": 0},
+        }
+    ]
+    assert len(_batch_writes(mock_prisma_client, "budget")) == 1
+
+
+_LINKED_TABLE_CASES = [
+    ("team_membership", {"budget_id": {"in": ["7d-budget-tier"]}}),
+    (
+        "key",
+        {
+            "budget_id": {"in": ["7d-budget-tier"]},
+            "budget_duration": None,
+            "spend": {"gt": 0},
         },
-    )
-
-    budgets_to_reset = [test_budget]
-
-    # Run the method
-    asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset))
-
-    # Verify that update_many was called on litellm_verificationtoken
-    calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls
-    assert len(calls) == 1, f"Expected 1 update_many call, got {len(calls)}"
-
-    # Verify the where clause filters by budget_id and null budget_duration
-    call = calls[0]
-    assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]}
-    assert call["where"]["budget_duration"] is None
-
-    # Verify spend is reset to 0
-    assert call["data"]["spend"] == 0
+    ),
+    ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}),
+    ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}),
+]
 
 
-def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_duration(
-    reset_budget_job, mock_prisma_client
+@pytest.mark.parametrize(
+    "table, expected_where",
+    _LINKED_TABLE_CASES,
+    ids=[case[0] for case in _LINKED_TABLE_CASES],
+)
+def test_budget_table_reset_zeroes_spend_on_every_linked_table(
+    reset_budget_job, mock_prisma_client, table, expected_where
 ):
+    """One expiring tier zeroes spend on every row it gates.
+
+    The filters carry real behavior: keys must be narrowed to
+    `budget_duration: None` so keys with their own reset schedule aren't
+    double-reset by reset_budget_for_litellm_keys(), and the payload must stay
+    exactly {"spend": 0} because `total_spend` is a lifetime counter a reset
+    may never touch.
     """
-    Test that keys with BOTH budget_id AND budget_duration are excluded from
-    reset_budget_for_keys_linked_to_budgets. Such keys have their own reset
-    schedule and are handled only by reset_budget_for_litellm_keys(). The
-    budget_duration=None filter ensures they are NOT double-reset when the
-    linked budget tier expires.
-    """
-    now = datetime.now(timezone.utc)
+    mock_prisma_client.data["budget"] = [_budget_row(budget_id="7d-budget-tier", budget_duration="7d")]
 
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "max_budget": 10.0,
-            "budget_duration": "7d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "7d-budget-tier",
-            "created_at": now - timedelta(days=7),
-        },
-    )
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    budgets_to_reset = [test_budget]
-
-    asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset))
-
-    calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls
-    assert len(calls) == 1
-    call = calls[0]
-
-    # Critical: budget_duration must be None so keys with their own budget_duration
-    # (e.g. key has budget_id="X" AND budget_duration=60) are excluded.
-    # Those keys are reset only by reset_budget_for_litellm_keys() - no double-reset.
-    assert call["where"]["budget_duration"] is None
-    assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]}
+    writes = _batch_writes(mock_prisma_client, table, op="update_many")
+    assert len(writes) == 1, f"expected exactly 1 {table} write, got {writes}"
+    assert writes[0]["where"] == expected_where
+    assert writes[0]["data"] == {"spend": 0}
 
 
-def test_reset_budget_for_keys_linked_to_budgets_empty(reset_budget_job, mock_prisma_client):
-    """
-    Test that when there are no budgets to reset, no update is performed
-    on the verification token table.
-    """
-    # Run with empty list
-    asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[]))
+def test_budget_table_reset_writes_nothing_when_no_budget_is_due(reset_budget_job, mock_prisma_client):
+    """Nothing due means no transaction is opened at all."""
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    # Verify no update_many calls were made
-    calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls
-    assert len(calls) == 0
+    assert mock_prisma_client.db.batchers == []
+    assert mock_prisma_client.db.batch_calls == []
 
 
-def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_client):
-    """
-    Test that when a budget tier is reset, orgs linked to that budget
-    (via budget_id) also get their spend reset.
-    """
-    now = datetime.now(timezone.utc)
+def _run_reset_at_fixed_now(job, fixed_now):
+    """Run the budget-table reset with `now` pinned for reset-time math."""
+    from unittest.mock import patch
 
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "max_budget": 100.0,
-            "budget_duration": "30d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "30d-org-budget",
-            "created_at": now - timedelta(days=30),
-        },
-    )
-
-    asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[test_budget]))
-
-    calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls
-    assert len(calls) == 1
-    call = calls[0]
-    assert call["where"]["budget_id"] == {"in": ["30d-org-budget"]}
-    assert call["where"]["spend"] == {"gt": 0}
-    assert call["data"]["spend"] == 0
-
-
-def test_reset_budget_for_orgs_linked_to_budgets_empty(reset_budget_job, mock_prisma_client):
-    """
-    Test that when there are no budgets to reset, no update is performed
-    on the organization table.
-    """
-    asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[]))
-    calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls
-    assert len(calls) == 0
-
-
-def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_client):
-    """
-    Test that when a budget tier is reset, tags linked to that budget
-    (via budget_id) also get their spend reset.
-    """
-    now = datetime.now(timezone.utc)
-
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "max_budget": 50.0,
-            "budget_duration": "30d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "30d-tag-budget",
-            "created_at": now - timedelta(days=30),
-        },
-    )
-
-    asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[test_budget]))
-
-    calls = mock_prisma_client.db.litellm_tagtable.update_many_calls
-    assert len(calls) == 1
-    call = calls[0]
-    assert call["where"]["budget_id"] == {"in": ["30d-tag-budget"]}
-    assert call["where"]["spend"] == {"gt": 0}
-    assert call["data"]["spend"] == 0
-
-
-def test_reset_budget_for_tags_linked_to_budgets_empty(reset_budget_job, mock_prisma_client):
-    """
-    Test that when there are no budgets to reset, no update is performed
-    on the tag table.
-    """
-    asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[]))
-    calls = mock_prisma_client.db.litellm_tagtable.update_many_calls
-    assert len(calls) == 0
+    with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt:
+        mock_dt.now.return_value = fixed_now
+        mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
+        asyncio.run(job.reset_budget_for_litellm_budget_table())
 
 
 @pytest.mark.parametrize(
@@ -680,215 +572,70 @@ def test_reset_budget_for_tags_linked_to_budgets_empty(reset_budget_job, mock_pr
     ],
     ids=["30d-calendar-month", "1mo-calendar-month", "1d-next-midnight"],
 )
-def test_reset_budget_reset_at_date_calendar_aligned(budget_duration, expected_day, expected_month):
-    """
-    Verify that _reset_budget_reset_at_date produces calendar-aligned reset
-    times (matching get_budget_reset_time), not sliding-window offsets.
-    """
-    from unittest.mock import patch
-
-    # Fix "now" to 2023-06-15 10:30:00 UTC for deterministic results
+def test_budget_reset_at_written_is_calendar_aligned(
+    reset_budget_job, mock_prisma_client, budget_duration, expected_day, expected_month
+):
+    """The advanced budget_reset_at is calendar-aligned, not a sliding
+    now + duration offset."""
     fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc)
+    mock_prisma_client.data["budget"] = [
+        _budget_row(
+            budget_id="test-budget",
+            budget_duration=budget_duration,
+            budget_reset_at=fixed_now - timedelta(hours=1),
+        )
+    ]
 
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "budget_duration": budget_duration,
-            "budget_reset_at": fixed_now - timedelta(hours=1),
-            "budget_id": "test-budget",
-            "created_at": fixed_now - timedelta(days=30),
-        },
-    )
+    _run_reset_at_fixed_now(reset_budget_job, fixed_now)
 
-    with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt:
-        mock_dt.now.return_value = fixed_now
-        mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
-        asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings()))
-
-    assert test_budget.budget_reset_at.day == expected_day
-    assert test_budget.budget_reset_at.month == expected_month
-    assert test_budget.budget_reset_at.hour == 0
-    assert test_budget.budget_reset_at.minute == 0
-    assert test_budget.budget_reset_at.second == 0
+    writes = _batch_writes(mock_prisma_client, "budget")
+    assert len(writes) == 1
+    written = writes[0]["data"]["budget_reset_at"]
+    assert (written.day, written.month) == (expected_day, expected_month)
+    assert (written.hour, written.minute, written.second) == (0, 0, 0)
 
 
-def test_reset_budget_reset_at_date_7d_next_monday():
-    """Verify 7d budget duration resets to next Monday at midnight."""
-    from unittest.mock import patch
-
+def test_budget_reset_at_written_for_7d_is_next_monday(reset_budget_job, mock_prisma_client):
+    """7d budgets advance to next Monday at midnight."""
     # 2023-06-14 is a Wednesday
     fixed_now = datetime(2023, 6, 14, 10, 30, 0, tzinfo=timezone.utc)
+    mock_prisma_client.data["budget"] = [
+        _budget_row(budget_id="test-budget", budget_duration="7d", budget_reset_at=fixed_now - timedelta(hours=1))
+    ]
 
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "budget_duration": "7d",
-            "budget_reset_at": fixed_now - timedelta(hours=1),
-            "budget_id": "test-budget",
-            "created_at": fixed_now - timedelta(days=7),
-        },
-    )
+    _run_reset_at_fixed_now(reset_budget_job, fixed_now)
 
-    with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt:
-        mock_dt.now.return_value = fixed_now
-        mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
-        asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings()))
-
-    # Next Monday after Wednesday June 14 is June 19
-    assert test_budget.budget_reset_at.day == 19
-    assert test_budget.budget_reset_at.month == 6
-    assert test_budget.budget_reset_at.weekday() == 0  # Monday
-    assert test_budget.budget_reset_at.hour == 0
+    written = _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"]
+    assert (written.day, written.month) == (19, 6)
+    assert written.weekday() == 0
+    assert written.hour == 0
 
 
-def test_reset_budget_reset_at_date_none_duration():
-    """Verify that budget_reset_at is unchanged when budget_duration is None."""
-    original_reset_at = datetime(2023, 6, 20, 0, 0, 0, tzinfo=timezone.utc)
-    now = datetime(2023, 6, 15, 10, 0, 0, tzinfo=timezone.utc)
+def test_budget_with_no_duration_gets_no_reset_at_write(reset_budget_job, mock_prisma_client):
+    """A tier without a duration has no next window, so its row is left alone
+    rather than rewritten with an unchanged value."""
+    mock_prisma_client.data["budget"] = [
+        _budget_row(
+            budget_id="no-duration", budget_duration=None, budget_reset_at=datetime(2023, 6, 20, tzinfo=timezone.utc)
+        )
+    ]
 
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "budget_duration": None,
-            "budget_reset_at": original_reset_at,
-            "budget_id": "test-budget",
-            "created_at": now - timedelta(days=30),
-        },
-    )
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now, BudgetResetSettings()))
-    assert test_budget.budget_reset_at == original_reset_at
+    assert _batch_writes(mock_prisma_client, "budget") == []
 
 
-def test_reset_budget_reset_at_date_none_reset_at():
-    """Verify that budget_reset_at is set correctly even when previously None."""
-    from unittest.mock import patch
-
+def test_budget_reset_at_written_when_previously_null(reset_budget_job, mock_prisma_client):
+    """A tier whose budget_reset_at was never initialized still gets one."""
     fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc)
+    mock_prisma_client.data["budget"] = [
+        _budget_row(budget_id="test-budget", budget_duration="30d", budget_reset_at=None)
+    ]
 
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "budget_duration": "30d",
-            "budget_reset_at": None,
-            "budget_id": "test-budget",
-            "created_at": fixed_now - timedelta(days=5),
-        },
-    )
+    _run_reset_at_fixed_now(reset_budget_job, fixed_now)
 
-    with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt:
-        mock_dt.now.return_value = fixed_now
-        mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs)
-        asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings()))
-
-    # Should be set to 1st of next month (July 1)
-    assert test_budget.budget_reset_at is not None
-    assert test_budget.budget_reset_at.day == 1
-    assert test_budget.budget_reset_at.month == 7
-
-
-def test_budget_table_reset_also_resets_linked_keys(reset_budget_job, mock_prisma_client):
-    """
-    Integration-style test: when reset_budget_for_litellm_budget_table runs,
-    it should also reset spend for keys linked to the expiring budget tiers
-    (in addition to end-users and team members).
-    """
-    now = datetime.now(timezone.utc)
-
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "max_budget": 10.0,
-            "budget_duration": "7d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "7d-budget-tier",
-            "created_at": now - timedelta(days=7),
-        },
-    )
-
-    mock_prisma_client.data["budget"] = [test_budget]
-
-    # Run the full budget table reset
-    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
-
-    # Verify that keys linked to the budget were also reset
-    calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls
-    assert len(calls) == 1, (
-        "Expected reset_budget_for_litellm_budget_table to also reset keys "
-        f"linked to expiring budgets, but got {len(calls)} update_many calls"
-    )
-    assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]}
-    assert calls[0]["data"]["spend"] == 0
-
-
-def test_budget_table_reset_also_resets_linked_orgs(reset_budget_job, mock_prisma_client):
-    """
-    Integration-style test: when reset_budget_for_litellm_budget_table runs,
-    it should also reset spend for orgs linked to the expiring budget tiers
-    (in addition to end-users, team members, and keys).
-    """
-    now = datetime.now(timezone.utc)
-
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "max_budget": 100.0,
-            "budget_duration": "30d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "30d-org-budget",
-            "created_at": now - timedelta(days=30),
-        },
-    )
-
-    mock_prisma_client.data["budget"] = [test_budget]
-
-    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
-
-    calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls
-    assert len(calls) == 1, (
-        "Expected reset_budget_for_litellm_budget_table to also reset orgs "
-        f"linked to expiring budgets, but got {len(calls)} update_many calls"
-    )
-    assert calls[0]["where"]["budget_id"] == {"in": ["30d-org-budget"]}
-    assert calls[0]["data"]["spend"] == 0
-
-
-def test_budget_table_reset_also_resets_linked_tags(reset_budget_job, mock_prisma_client):
-    """
-    Integration-style test: when reset_budget_for_litellm_budget_table runs,
-    it should also reset spend for tags linked to the expiring budget tiers.
-    """
-    now = datetime.now(timezone.utc)
-
-    test_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {
-            "max_budget": 50.0,
-            "budget_duration": "30d",
-            "budget_reset_at": now - timedelta(hours=1),
-            "budget_id": "30d-tag-budget",
-            "created_at": now - timedelta(days=30),
-        },
-    )
-
-    mock_prisma_client.data["budget"] = [test_budget]
-
-    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
-
-    calls = mock_prisma_client.db.litellm_tagtable.update_many_calls
-    assert len(calls) == 1, (
-        "Expected reset_budget_for_litellm_budget_table to also reset tags "
-        f"linked to expiring budgets, but got {len(calls)} update_many calls"
-    )
-    assert calls[0]["where"]["budget_id"] == {"in": ["30d-tag-budget"]}
-    assert calls[0]["data"]["spend"] == 0
+    written = _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"]
+    assert (written.day, written.month) == (1, 7)
 
 
 def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock_prisma_client):
@@ -965,16 +712,14 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock
 
     asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    # Both end users should have been reset
-    updated = mock_prisma_client.updated_data["enduser"]
-    assert len(updated) == 2, f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}"
-
-    user_ids = {u.user_id for u in updated}
-    assert "enduser-explicit" in user_ids
-    assert "enduser-implicit" in user_ids
-
-    for u in updated:
-        assert u.spend == 0.0, f"Expected spend=0 for {u.user_id}, got {u.spend}"
+    # Both end users are zeroed by the same committed statement.
+    enduser_writes = _batch_writes(mock_prisma_client, "enduser")
+    assert len(enduser_writes) == 1, f"Expected a single enduser write, got {enduser_writes}"
+    assert set(enduser_writes[0]["where"]["user_id"]["in"]) == {
+        "enduser-explicit",
+        "enduser-implicit",
+    }
+    assert enduser_writes[0]["data"] == {"spend": 0}
 
     # Verify find_many was called to fetch NULL-budget-id end users
     find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
@@ -1054,34 +799,6 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li
     litellm.max_end_user_budget_id = None
 
 
-def test_reset_budget_for_team_members_preserves_total_spend():
-    """Regression guard: reset_budget_for_litellm_team_members must zero `spend`
-    but leave `total_spend` untouched.
-
-    The reset writes `data={"spend": 0}` explicitly. If a future refactor adds
-    `"total_spend": 0` to that dict, this test fails immediately.
-    """
-    expired_budget = type(
-        "LiteLLM_BudgetTableFull",
-        (),
-        {"budget_id": "budget-1"},
-    )
-
-    mock_prisma_client = MagicMock()
-    mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[])
-    mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client)
-
-    asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
-
-    mock_prisma_client.db.litellm_teammembership.update_many.assert_called_once()
-    call_kwargs = mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs
-    assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"]
-    assert call_kwargs["data"] == {"spend": 0}
-    assert "total_spend" not in call_kwargs["data"]
-
-
 # ---------------------------------------------------------------------------
 # reset_budget_windows (per-key / per-team concurrent window resets)
 # ---------------------------------------------------------------------------
@@ -1323,28 +1040,6 @@ def _make_counter_invalidation_job(monkeypatch):
     return spend_counter_cache
 
 
-def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch):
-    """Team-member budget reset clears the Redis spend counter."""
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    membership = type(
-        "Membership",
-        (),
-        {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
-    )
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership])
-    prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
-
-    counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team_member:alice:team-x", value=0.0, ttl=60)
-    counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:team_member:alice:team-x", value=0.0, ttl=60)
-
-
 def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch):
     """Key budget reset must clear the Redis spend counter."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)
@@ -1574,207 +1269,240 @@ def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job,
     )
 
 
-def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch):
-    """Resetting keys via budget tier must clear each linked key's counter."""
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_key = type("Key", (), {"token": "sk-linked"})
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key])
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
-
-    counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-linked", value=0.0, ttl=60)
+_INVALIDATION_CASES = [
+    (
+        "litellm_teammembership",
+        type("Membership", (), {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}),
+        "spend:team_member:alice:team-x",
+        {"team-x_alice"},
+    ),
+    (
+        "litellm_verificationtoken",
+        type("Key", (), {"token": "sk-linked"}),
+        "spend:key:sk-linked",
+        {"sk-linked"},
+    ),
+    (
+        "litellm_organizationtable",
+        type("Org", (), {"organization_id": "org-acme"}),
+        "spend:org:org-acme",
+        {"org_id:org-acme", "org_id:org-acme:with_budget"},
+    ),
+    (
+        "litellm_tagtable",
+        type("Tag", (), {"tag_name": "tenant-42"}),
+        "spend:tag:tenant-42",
+        {"tag:tenant-42"},
+    ),
+]
 
 
-def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monkeypatch):
-    """Resetting orgs via budget tier must clear each linked org's counter."""
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_org = type("Org", (), {"organization_id": "org-acme"})
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org])
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget]))
-
-    counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:org:org-acme", value=0.0, ttl=60)
-    counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:org:org-acme", value=0.0, ttl=60)
-
-
-def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monkeypatch):
-    """Resetting tags via budget tier must clear each linked tag's counter."""
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_tag = type("Tag", (), {"tag_name": "tenant-42"})
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag])
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
-
-    counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:tag:tenant-42", value=0.0, ttl=60)
-    counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:tag:tenant-42", value=0.0, ttl=60)
-
-
-def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache(
-    monkeypatch,
+@pytest.mark.parametrize(
+    "table_attr, linked_row, counter_key, cache_keys",
+    _INVALIDATION_CASES,
+    ids=["team_membership", "key", "org", "tag"],
+)
+def test_budget_table_reset_invalidates_counters_and_management_cache(
+    reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys
 ):
-    """Regression guard for the bug where tag spend stayed frozen across cycles.
+    """Every row the cascade zeroes gets its spend counter cleared and its
+    management-cache entry dropped.
 
-    ``SpendCounterReseed.from_db`` returns ``None`` for ``spend:tag:*`` keys,
-    so once the spend counter expires the tag budget check falls back to the
-    cached ``LiteLLM_TagTable.spend``. If we don't drop the management cache
-    entry on reset, that cached object lingers (TTL 60s) with the pre-reset
-    spend, and ``_tag_max_budget_check`` keeps returning HTTP 400 even though
-    the DB row has been zeroed.
+    Both matter. ``SpendCounterReseed.from_db`` returns None for tags, so once
+    the counter expires the budget check falls back to the cached row's
+    ``.spend``; and for keys, orgs and team memberships another pod's cached
+    object can stay pinned above the zeroed DB row until its TTL. Team
+    membership cache keys follow auth's ``{team_id}_{user_id}`` shape, and orgs
+    carry both the plain and the ``:with_budget`` entry.
     """
     counter_cache = _make_counter_invalidation_job(monkeypatch)
+    mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
+    getattr(mock_prisma_client.db, table_attr).set_find_many_results([linked_row])
 
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_tag = type("Tag", (), {"tag_name": "tenant-42"})
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag])
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
-
-    counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="tag:tenant-42")
+    counter_cache.in_memory_cache.set_cache.assert_any_call(key=counter_key, value=0.0, ttl=60)
+    counter_cache.redis_cache.async_set_cache.assert_any_await(key=counter_key, value=0.0, ttl=60)
+    deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
+    assert cache_keys <= deleted
 
 
-def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management_cache(
-    monkeypatch,
-):
-    """When multiple tags share the expired budget tier, every one of them
-    has its ``user_api_key_cache`` entry dropped — not just the first."""
+def test_budget_table_reset_invalidates_every_tag_not_just_the_first(reset_budget_job, mock_prisma_client, monkeypatch):
+    """When several tags share the expiring tier, all of them are evicted."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_tags = [
-        type("Tag", (), {"tag_name": "tenant-a"}),
-        type("Tag", (), {"tag_name": "tenant-b"}),
-        type("Tag", (), {"tag_name": "tenant-c"}),
-    ]
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=linked_tags)
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 3})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
-
-    deleted_keys = {
-        call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list
-    }
-    assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"}
-
-
-def test_reset_budget_for_keys_linked_to_budgets_invalidates_management_cache(
-    monkeypatch,
-):
-    """Budget-tier key resets must drop the cached key object (hashed token key).
-
-    Historically this test used ``assert_not_awaited()`` on
-    ``user_api_key_cache.async_delete_cache``, reflecting the assumption that
-    ``SpendCounterReseed.from_db`` alone kept spend consistent for keys and
-    that invalidating the management cache was unnecessary. That was flipped to
-    ``assert_any_await(...)`` because the old invariant fails across pods: a
-    budget reset on one instance can leave another pod's cached key object
-    (including embedded ``.spend``) stale until TTL expiry. Eviction now matches
-    tags/orgs/teams. Do not treat the ``cache_key_fn`` / invalidation wiring as
-    redundant without revisiting that cross-pod consistency story.
-    """
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_key = type("Key", (), {"token": "sk-linked"})
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key])
-    prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
-
-    counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="sk-linked")
-
-
-def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache(
-    monkeypatch,
-):
-    """Org rows use both base and budget-table cache keys — evict both on reset."""
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_org = type("Org", (), {"organization_id": "org-acme"})
-
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org])
-    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1})
-
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget]))
-
-    deleted_keys = {
-        call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list
-    }
-    assert deleted_keys == {
-        "org_id:org-acme",
-        "org_id:org-acme:with_budget",
-    }
-
-
-def test_reset_budget_for_team_members_invalidates_management_cache(monkeypatch):
-    """Team membership cache key matches auth: ``{team_id}_{user_id}``."""
-    counter_cache = _make_counter_invalidation_job(monkeypatch)
-
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    membership = type(
-        "Membership",
-        (),
-        {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
+    mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
+    mock_prisma_client.db.litellm_tagtable.set_find_many_results(
+        [type("Tag", (), {"tag_name": name}) for name in ("tenant-a", "tenant-b", "tenant-c")]
     )
 
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership])
-    prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1})
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
-
-    counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="team-x_alice")
+    deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list}
+    assert deleted == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"}
 
 
-def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure_still_resets(
-    monkeypatch,
-):
-    """If ``async_delete_cache`` raises, the DB cascade must still complete."""
+def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch):
+    """Eviction runs after the commit, so a broken cache cannot undo the write."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)
     counter_cache.user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("cache unavailable"))
+    mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
+    mock_prisma_client.db.litellm_tagtable.set_find_many_results([type("Tag", (), {"tag_name": "tenant-42"})])
 
-    expired_budget = type("B", (), {"budget_id": "budget-1"})
-    linked_tag = type("Tag", (), {"tag_name": "tenant-42"})
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
 
-    prisma_client = MagicMock()
-    prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag])
-    prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1})
+    assert len(_batch_writes(mock_prisma_client, "tag", op="update_many")) == 1
+    assert mock_prisma_client.db.batchers[0].committed is True
 
-    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
-    asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
 
-    prisma_client.db.litellm_tagtable.update_many.assert_awaited_once()
+# ---------------------------------------------------------------------------
+# Atomicity of the budget-table cascade (LIT-5138)
+# ---------------------------------------------------------------------------
+
+
+class FailingCommitDB(MockDB):
+    """Batches that blow up at commit, like a Postgres timeout mid-cascade."""
+
+    def batch_(self):
+        batcher = super().batch_()
+
+        async def _fail():
+            raise RuntimeError("simulated Postgres timeout mid-cascade")
+
+        batcher.commit = _fail
+        return batcher
+
+
+class FailingTeamMembershipDB(MockDB):
+    """Queueing the team-membership reset raises, i.e. the cascade breaks after
+    earlier writes are already queued."""
+
+    def batch_(self):
+        batcher = super().batch_()
+
+        def _fail(where, data):
+            raise RuntimeError("simulated failure queueing the team-membership reset")
+
+        batcher.litellm_teammembership.update_many = _fail
+        return batcher
+
+
+class OrderRecordingDB(MockDB):
+    """Appends a marker to a shared list when a batch commits."""
+
+    def __init__(self, events):
+        super().__init__()
+        self._events = events
+
+    def batch_(self):
+        batcher = super().batch_()
+        wrapped = batcher.commit
+
+        async def _record_commit():
+            self._events.append("commit")
+            return await wrapped()
+
+        batcher.commit = _record_commit
+        return batcher
+
+
+def _job_with_expired_budget(db, proxy_logging=None):
+    """A job with one due tier and a linked tag, so cache invalidation has
+    something to invalidate and its absence is a real signal."""
+    prisma_client = MockPrismaClient()
+    prisma_client.db = db
+    prisma_client.data["budget"] = [_budget_row(budget_id="budget-1", budget_duration="7d")]
+    db.litellm_tagtable.set_find_many_results([type("Tag", (), {"tag_name": "tenant-42"})])
+    job = ResetBudgetJob(
+        proxy_logging_obj=proxy_logging or MockProxyLogging(),
+        prisma_client=prisma_client,
+    )
+    return job, prisma_client
+
+
+@pytest.mark.parametrize(
+    "db_factory",
+    [FailingCommitDB, FailingTeamMembershipDB],
+    ids=["commit-fails", "queueing-fails"],
+)
+def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monkeypatch):
+    """Regression for LIT-5138.
+
+    The old code committed the new budget_reset_at first and zeroed the
+    dependent spend afterwards. A failure part-way through left the tier
+    stamped for the next window, so every later tick skipped it and team
+    member / enduser / org / tag spend stayed at the cap for the whole window.
+    One transaction means a failure anywhere persists nothing and the tier is
+    still due on the next tick.
+    """
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+    job, prisma_client = _job_with_expired_budget(db_factory())
+
+    asyncio.run(job.reset_budget_for_litellm_budget_table())  # swallowed, retried next tick
+
+    assert prisma_client.db.batch_calls == [], "a failed cascade must not persist any write"
+    assert prisma_client.db.batchers[0].committed is False
+    assert prisma_client.updated_data["budget"] == [], "budget_reset_at must not be advanced outside the transaction"
+    counter_cache.in_memory_cache.set_cache.assert_not_called()
+    counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited()
+
+
+def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mock_prisma_client, monkeypatch):
+    """Dependent spend and the budget_reset_at advance ride one batch."""
+    _make_counter_invalidation_job(monkeypatch)
+    now = datetime.now(timezone.utc)
+    budget = _budget_row(budget_id="budget-1", budget_duration="7d")
+    mock_prisma_client.data["budget"] = [budget]
+    mock_prisma_client.data["enduser"] = [
+        type("EndUser", (), {"spend": 5.0, "litellm_budget_table": budget, "user_id": "enduser-1"})
+    ]
+
+    asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
+
+    assert len(mock_prisma_client.db.batchers) == 1, "the cascade must not be split across transactions"
+    batcher = mock_prisma_client.db.batchers[0]
+    assert batcher.committed is True
+    assert {(call["table"], call["op"]) for call in batcher.calls} == {
+        ("team_membership", "update_many"),
+        ("key", "update_many"),
+        ("org", "update_many"),
+        ("tag", "update_many"),
+        ("enduser", "update_many"),
+        ("budget", "update_many"),
+    }
+    budget_write = next(call for call in batcher.calls if call["table"] == "budget")
+    assert budget_write["data"]["budget_reset_at"] > now
+
+
+def test_caches_are_invalidated_only_after_the_transaction_commits(monkeypatch):
+    """A counter zeroed before the write lands would admit requests past the
+    cap while the DB still holds the over-budget spend."""
+    events = []
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+    counter_cache.in_memory_cache.set_cache.side_effect = lambda **kwargs: events.append("counter")
+
+    job, _ = _job_with_expired_budget(OrderRecordingDB(events))
+
+    asyncio.run(job.reset_budget_for_litellm_budget_table())
+
+    assert events == ["commit", "counter"]
+
+
+def test_failed_cascade_is_logged_as_a_cascade_failure(monkeypatch):
+    """The failure log has to name what actually broke. The old catch-all
+    blamed end users even when the team-membership write was the failure."""
+    from unittest.mock import patch
+
+    _make_counter_invalidation_job(monkeypatch)
+    job, _ = _job_with_expired_budget(FailingTeamMembershipDB())
+
+    with patch("litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception") as mock_exception:
+        asyncio.run(job.reset_budget_for_litellm_budget_table())
+
+    assert mock_exception.call_count == 1
+    message = mock_exception.call_args.args[0]
+    assert "cascade" in message
+    for mentioned in ("team member", "enduser", "org", "tag", "budget_reset_at"):
+        assert mentioned in message, f"failure log should mention {mentioned}: {message}"
 
 
 def _extract_reset_where(find_many_mock):
@@ -1799,23 +1527,65 @@ def _asserts_null_reset_is_due(where):
     branches = where.get("OR")
     assert isinstance(branches, list), f"expected an OR filter, got {where!r}"
 
-    has_null_branch = any(
-        b.get("AND")
-        == [
-            {"budget_reset_at": None},
-            {"NOT": {"budget_duration": None}},
-        ]
-        for b in branches
-        if isinstance(b, dict)
-    )
-    has_expired_branch = any(
-        isinstance(b, dict)
-        and "budget_reset_at" in b
-        and b["budget_reset_at"] is not None
-        for b in branches
-    )
+    has_null_branch = {"budget_reset_at": None} in branches
+    has_expired_branch = any(isinstance(b, dict) and isinstance(b.get("budget_reset_at"), dict) for b in branches)
     assert has_null_branch, f"missing NULL-reset_at branch in {where!r}"
     assert has_expired_branch, f"missing expired-reset_at branch in {where!r}"
+    assert where.get("NOT") == {"budget_duration": None}, f"NULL reset_at is only due with a duration: {where!r}"
+
+
+_RESET_TABLE_ATTRS = {
+    "user": "litellm_usertable",
+    "team": "litellm_teamtable",
+    "budget": "litellm_budgettable",
+    "key": "litellm_verificationtoken",
+}
+
+
+def _run_reset_query(table_name, **extra):
+    """Run ``get_data`` for one table's budget-reset query against a mocked
+    prisma handle, and hand back the ``find_many`` mock it drove."""
+    from litellm.proxy.utils import PrismaClient
+
+    client = PrismaClient.__new__(PrismaClient)
+    client.db = MagicMock()
+    find_many = AsyncMock(return_value=[])
+    setattr(getattr(client.db, _RESET_TABLE_ATTRS[table_name]), "find_many", find_many)
+
+    now = datetime.now(timezone.utc)
+    expires = {"expires": now} if table_name == "key" else {}
+    asyncio.run(client.get_data(table_name=table_name, query_type="find_all", reset_at=now, **expires, **extra))
+    return find_many
+
+
+@pytest.mark.parametrize("table_name", ["user", "team", "budget", "key"])
+def test_get_data_reset_query_applies_the_row_limit(table_name):
+    """The reset job pages through due rows, so ``limit`` has to reach prisma as
+    ``take``. Dropped, every worker goes back to pulling the entire expired set
+    in one unbounded query at the same calendar boundary."""
+    find_many = _run_reset_query(table_name, limit=7)
+
+    assert find_many.await_args.kwargs["take"] == 7
+
+
+@pytest.mark.parametrize("table_name", ["user", "team", "budget", "key"])
+def test_get_data_reset_query_skips_rows_with_no_budget_duration(table_name):
+    """A row with a past budget_reset_at but no budget_duration has no next
+    window to move to, so it stays due forever. Fetching it means re-reading and
+    re-zeroing it on every tick, and a full chunk of such rows makes the paged
+    scan report no progress and starve the whole phase.
+    """
+    find_many = _run_reset_query(table_name)
+
+    assert find_many.await_args.kwargs["where"]["NOT"] == {"budget_duration": None}
+
+
+@pytest.mark.parametrize("table_name", ["user", "team", "budget", "key"])
+def test_get_data_reset_query_is_unlimited_when_no_limit_is_passed(table_name):
+    """Callers that pass no limit keep the old unbounded behaviour."""
+    find_many = _run_reset_query(table_name)
+
+    assert find_many.await_args.kwargs.get("take") is None
 
 
 @pytest.mark.parametrize("table_name", ["user", "team"])
@@ -1838,8 +1608,327 @@ def test_get_data_reset_query_selects_null_budget_reset_at(table_name):
     setattr(getattr(client.db, table_attr), "find_many", find_many)
 
     now = datetime.now(timezone.utc)
-    asyncio.run(
-        client.get_data(table_name=table_name, query_type="find_all", reset_at=now)
-    )
+    asyncio.run(client.get_data(table_name=table_name, query_type="find_all", reset_at=now))
 
     _asserts_null_reset_is_due(_extract_reset_where(find_many))
+
+
+def _key_row(token: str, budget_duration: Any = "30d"):
+    """A key that is already due for a reset, shaped like a get_data() row."""
+    now = datetime.now(timezone.utc)
+    return type(
+        "LiteLLM_VerificationToken",
+        (),
+        {
+            "spend": 100.0,
+            "budget_duration": budget_duration,
+            "budget_reset_at": now - timedelta(hours=1),
+            "token": token,
+        },
+    )
+
+
+def _user_row(user_id: str, budget_duration: Any = "30d"):
+    now = datetime.now(timezone.utc)
+    return type(
+        "LiteLLM_UserTable",
+        (),
+        {
+            "spend": 100.0,
+            "budget_duration": budget_duration,
+            "budget_reset_at": now - timedelta(hours=1),
+            "user_id": user_id,
+        },
+    )
+
+
+def _team_row(team_id: str, budget_duration: Any = "30d"):
+    now = datetime.now(timezone.utc)
+    return type(
+        "LiteLLM_TeamTable",
+        (),
+        {
+            "spend": 100.0,
+            "budget_duration": budget_duration,
+            "budget_reset_at": now - timedelta(hours=1),
+            "team_id": team_id,
+        },
+    )
+
+
+# ---------------------------------------------------------------------------
+# Chunked batches
+# ---------------------------------------------------------------------------
+
+
+class ChunkedPrismaClient(MockPrismaClient):
+    """Replays a scripted sequence of get_data chunks per table.
+
+    The last chunk repeats forever, so a phase that fails to terminate keeps
+    seeing rows rather than quietly running out of data.
+    """
+
+    def __init__(self, chunks_by_table: Dict[str, List[List[Any]]]):
+        super().__init__()
+        self._chunks_by_table = chunks_by_table
+        self.fetches_by_table: Dict[str, int] = {}
+
+    async def get_data(self, table_name, query_type, **kwargs):
+        self.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs})
+        chunks = self._chunks_by_table.get(table_name)
+        if not chunks:
+            return []
+        index = self.fetches_by_table.get(table_name, 0)
+        self.fetches_by_table[table_name] = index + 1
+        return chunks[min(index, len(chunks) - 1)]
+
+
+def _chunked_job(chunks_by_table):
+    client = ChunkedPrismaClient(chunks_by_table)
+    return client, ResetBudgetJob(proxy_logging_obj=MockProxyLogging(), prisma_client=client)
+
+
+def _fetch_limits(client, table_name):
+    return [call.get("limit") for call in client.get_data_calls if call["table_name"] == table_name]
+
+
+def test_key_reset_walks_the_due_rows_one_chunk_at_a_time(monkeypatch):
+    """Each chunk is fetched under a LIMIT and committed on its own batch, so a
+    large backlog never becomes one giant transaction."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client, job = _chunked_job({"key": [[_key_row("k1"), _key_row("k2")], [_key_row("k3")]]})
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    assert client.fetches_by_table["key"] == 2
+    assert _fetch_limits(client, "key") == [2, 2]
+    assert len(client.db.batchers) == 2
+    assert all(batcher.committed for batcher in client.db.batchers)
+    assert [len(batcher.calls) for batcher in client.db.batchers] == [2, 1]
+    assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2", "k3"]
+
+
+def test_key_reset_stops_after_a_chunk_shorter_than_the_batch_size(monkeypatch):
+    """Fewer rows than the limit means the backlog is drained, so no follow-up
+    query is worth issuing."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 5)
+    client, job = _chunked_job({"key": [[_key_row("k1")]]})
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    assert client.fetches_by_table["key"] == 1
+
+
+def test_key_reset_stops_when_a_full_chunk_advances_nothing(monkeypatch):
+    """A key with no budget_duration keeps its past budget_reset_at, so the very
+    same rows come back on the next fetch. Treating those writes as progress
+    would re-read that chunk until the iteration cap, every tick."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    stuck_chunk = [_key_row("k1", budget_duration=None), _key_row("k2", budget_duration=None)]
+    client, job = _chunked_job({"key": [stuck_chunk]})
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    assert client.fetches_by_table["key"] == 1
+
+
+def test_key_reset_stops_when_the_fetch_fails(monkeypatch):
+    """A phase whose query raises has made no progress; retrying it in a tight
+    loop would just hammer a struggling database."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client, job = _chunked_job({"key": [[_key_row("k1"), _key_row("k2")]]})
+
+    async def _boom(table_name, query_type, **kwargs):
+        client.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs})
+        raise RuntimeError("db is down")
+
+    client.get_data = _boom
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    assert len(client.get_data_calls) == 1
+
+
+def test_key_reset_is_capped_at_max_chunks_per_run(monkeypatch):
+    """Backstop against a phase that keeps making progress forever: the run ends
+    and the leftovers wait for the next tick."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 1)
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", 3)
+    client, job = _chunked_job({"key": [[_key_row("k1")]]})
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    assert client.fetches_by_table["key"] == 3
+
+
+@pytest.mark.parametrize(
+    "phase, table_name, row_factory",
+    [
+        ("reset_budget_for_litellm_users", "user", lambda uid: _user_row(uid)),
+        ("reset_budget_for_litellm_teams", "team", lambda tid: _team_row(tid)),
+    ],
+    ids=["users", "teams"],
+)
+def test_user_and_team_resets_are_chunked_too(monkeypatch, phase, table_name, row_factory):
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client, job = _chunked_job({table_name: [[row_factory("a"), row_factory("b")], [row_factory("c")]]})
+
+    asyncio.run(getattr(job, phase)())
+
+    assert client.fetches_by_table[table_name] == 2
+    assert _fetch_limits(client, table_name) == [2, 2]
+    assert len(client.db.batchers) == 2
+    assert len(_batch_writes(client, table_name, op="update")) == 3
+
+
+def test_budget_table_reset_walks_chunks_until_it_runs_dry(monkeypatch):
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client, job = _chunked_job({"budget": [[_budget_row("b1"), _budget_row("b2")], [_budget_row("b3")]]})
+
+    asyncio.run(job.reset_budget_for_litellm_budget_table())
+
+    assert client.fetches_by_table["budget"] == 2
+    assert _fetch_limits(client, "budget") == [2, 2]
+    assert len(client.db.batchers) == 2
+    assert all(batcher.committed for batcher in client.db.batchers)
+    assert [w["where"]["budget_id"] for w in _batch_writes(client, "budget", op="update_many")] == ["b1", "b2", "b3"]
+
+
+def test_budget_table_reset_stops_when_a_full_chunk_advances_no_window(monkeypatch):
+    """A tier with no budget_duration has its linked spend zeroed but keeps its
+    past budget_reset_at, so it stays due. Counting those spend writes as
+    progress would re-read the same chunk until the cap."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    stuck_chunk = [_budget_row("b1", budget_duration=None), _budget_row("b2", budget_duration=None)]
+    client, job = _chunked_job({"budget": [stuck_chunk]})
+
+    asyncio.run(job.reset_budget_for_litellm_budget_table())
+
+    assert client.fetches_by_table["budget"] == 1
+    assert _batch_writes(client, "budget", op="update_many") == []
+
+
+def test_budget_table_reset_stops_when_the_cascade_fails(monkeypatch):
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client, job = _chunked_job({"budget": [[_budget_row("b1"), _budget_row("b2")]]})
+    client.db = FailingCommitDB()
+
+    asyncio.run(job.reset_budget_for_litellm_budget_table())
+
+    assert client.fetches_by_table["budget"] == 1
+
+
+# ---------------------------------------------------------------------------
+# Progress means "no longer due", not "was written"
+# ---------------------------------------------------------------------------
+
+
+def test_key_reset_stops_when_the_new_reset_time_is_not_in_the_future(monkeypatch):
+    """A "0s" budget_duration resolves to the current time, so the row is written
+    and comes straight back on the next fetch. Treating a written row as progress
+    burns the whole per-run chunk cap on rows that never move.
+    """
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    stuck_chunk = [_key_row("k1", budget_duration="0s"), _key_row("k2", budget_duration="0s")]
+    client, job = _chunked_job({"key": [stuck_chunk]})
+
+    asyncio.run(job.reset_budget_for_litellm_keys())
+
+    assert client.fetches_by_table["key"] == 1
+    assert len(_batch_writes(client, "key", op="update")) == 2
+
+
+def test_budget_table_reset_stops_when_the_new_window_is_not_in_the_future(monkeypatch):
+    """Same zero-length window on the budget tier: advancing it to now leaves it
+    due, so the cascade must not report progress."""
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    stuck_chunk = [_budget_row("b1", budget_duration="0s"), _budget_row("b2", budget_duration="0s")]
+    client, job = _chunked_job({"budget": [stuck_chunk]})
+
+    asyncio.run(job.reset_budget_for_litellm_budget_table())
+
+    assert client.fetches_by_table["budget"] == 1
+    assert len(_batch_writes(client, "budget", op="update_many")) == 2
+
+
+class PoisonRow:
+    """A row the in-memory reset cannot write, like the DataError rows in #27730."""
+
+    token = "poison"
+    budget_duration = "30d"
+    budget_reset_at = None
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        raise RuntimeError("simulated failure resetting this row")
+
+
+class RecordingServiceLogging:
+    def __init__(self):
+        self.success_calls: List[Dict[str, Any]] = []
+        self.failure_calls: List[Dict[str, Any]] = []
+
+    async def async_service_success_hook(self, **kwargs):
+        self.success_calls.append(kwargs)
+
+    async def async_service_failure_hook(self, **kwargs):
+        self.failure_calls.append(kwargs)
+
+
+class RecordingProxyLogging:
+    def __init__(self):
+        self.service_logging_obj = RecordingServiceLogging()
+
+
+def _run_and_drain_hooks(make_coro):
+    """The service hooks are fired as tasks; give them a turn before asserting."""
+
+    async def _run():
+        await make_coro()
+        await asyncio.sleep(0.05)
+
+    asyncio.run(_run())
+
+
+def test_key_reset_keeps_paging_when_some_rows_in_a_chunk_fail(monkeypatch):
+    """One row that cannot be reset must not cost the phase its remaining chunks:
+    the rows that did reset are committed and are real progress, and the failure
+    is reported instead of aborting the run.
+    """
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client = ChunkedPrismaClient({"key": [[PoisonRow(), _key_row("k1")], [_key_row("k2")]]})
+    logging_obj = RecordingProxyLogging()
+    job = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=client)
+
+    _run_and_drain_hooks(job.reset_budget_for_litellm_keys)
+
+    assert client.fetches_by_table["key"] == 2
+    assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2"]
+    assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == ["reset_budget_keys"]
+    assert set(logging_obj.service_logging_obj.failure_calls[0]["event_metadata"]) == {
+        "num_keys_found",
+        "keys_found",
+    }
+    assert [call["call_type"] for call in logging_obj.service_logging_obj.success_calls] == ["reset_budget_keys"]
+
+
+@pytest.mark.parametrize(
+    "phase, table_name, row_factory, call_type",
+    [
+        ("reset_budget_for_litellm_users", "user", _user_row, "reset_budget_users"),
+        ("reset_budget_for_litellm_teams", "team", _team_row, "reset_budget_teams"),
+    ],
+    ids=["users", "teams"],
+)
+def test_user_and_team_chunks_report_progress_despite_a_failed_row(
+    monkeypatch, phase, table_name, row_factory, call_type
+):
+    monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2)
+    client = ChunkedPrismaClient({table_name: [[PoisonRow(), row_factory("a")], [row_factory("b")]]})
+    logging_obj = RecordingProxyLogging()
+    job = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=client)
+
+    _run_and_drain_hooks(getattr(job, phase))
+
+    assert client.fetches_by_table[table_name] == 2
+    assert len(_batch_writes(client, table_name, op="update")) == 2
+    assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == [call_type]
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 3bdf9bafdc7..6a9e894feb5 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py
@@ -72,6 +72,39 @@ async def test_new_budget_success(client_and_mocks):
     mock_table.create.assert_awaited_once()
 
 
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+@pytest.mark.asyncio
+async def test_new_budget_rejects_a_duration_that_never_advances(
+    client_and_mocks, bad_duration
+):
+    """A zero-length window resets to "now", so the row is due again the moment
+    it is written and the reset job re-reads it on every tick forever."""
+    client, _, mock_table = client_and_mocks
+
+    resp = client.post(
+        "/budget/new",
+        json={"budget_id": "budget_bad", "max_budget": 10.0, "budget_duration": bad_duration},
+    )
+
+    assert resp.status_code == 400, resp.text
+    assert "Invalid budget_duration" in resp.json()["detail"]["error"]
+    mock_table.create.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_update_budget_rejects_a_duration_that_never_advances(client_and_mocks):
+    client, _, mock_table = client_and_mocks
+
+    resp = client.post(
+        "/budget/update",
+        json={"budget_id": "budget_456", "budget_duration": "0s"},
+    )
+
+    assert resp.status_code == 400, resp.text
+    assert "Invalid budget_duration" in resp.json()["detail"]["error"]
+    mock_table.update.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_new_budget_db_not_connected(client_and_mocks, monkeypatch):
     client, mock_prisma, mock_table = client_and_mocks
diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py
index 81840745d0e..7dfd99dfa53 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py
@@ -628,6 +628,58 @@ class TestValidateFiniteSpendErrorDetail:
         }
 
 
+class TestValidateBudgetDuration:
+    """`validate_budget_duration` keeps durations that never advance out of the
+    database.
+
+    A duration of "0s" resolves to a reset time of now, so the row is due again
+    the instant it is written. The reset job re-reads such rows on every tick
+    and, once one tenant owns enough of them, they fill each batch and starve
+    every other tenant's reset.
+    """
+
+    def test_none_is_allowed(self):
+        from litellm.proxy.management_endpoints.common_utils import (
+            validate_budget_duration,
+        )
+
+        assert validate_budget_duration(None) is None
+
+    @pytest.mark.parametrize("duration", ["30s", "5m", "1h", "1d", "7d", "30d", "1mo"])
+    def test_positive_durations_are_allowed(self, duration):
+        from litellm.proxy.management_endpoints.common_utils import (
+            validate_budget_duration,
+        )
+
+        assert validate_budget_duration(duration) is None
+
+    @pytest.mark.parametrize("duration", ["0s", "0m", "0h", "0d", "-5m", "abc", ""])
+    def test_non_advancing_durations_are_rejected(self, duration):
+        from fastapi import HTTPException
+
+        from litellm.proxy.management_endpoints.common_utils import (
+            validate_budget_duration,
+        )
+
+        with pytest.raises(HTTPException) as exc_info:
+            validate_budget_duration(duration)
+        assert exc_info.value.status_code == 400
+
+    def test_rejection_detail_is_exact(self):
+        from fastapi import HTTPException
+
+        from litellm.proxy.management_endpoints.common_utils import (
+            validate_budget_duration,
+        )
+
+        with pytest.raises(HTTPException) as exc_info:
+            validate_budget_duration("0s")
+
+        assert exc_info.value.detail == {
+            "error": "Invalid budget_duration '0s'. Use a format like '1h', '24h', '7d', or '30d'."
+        }
+
+
 class TestRequireCallerUserIdErrorDetail:
     """The 403 for a service-account key must carry the exact error body."""
 
diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
index 0af5ad6cd9b..5efed8de325 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py
@@ -749,6 +749,40 @@ def test_char_new_body(mock_prisma_client, mock_user_api_key_auth):
     assert response.json() == _EXPECTED_CUSTOMER
 
 
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+def test_customer_new_rejects_a_duration_that_never_advances(
+    mock_prisma_client, mock_user_api_key_auth, bad_duration
+):
+    """A zero-length window resets to "now", leaving the customer's budget row
+    permanently due for the reset job to re-read every tick."""
+    mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW))
+
+    response = client.post(
+        "/customer/new",
+        json={"user_id": "c1", "max_budget": 10.0, "budget_duration": bad_duration},
+        headers={"Authorization": "Bearer k"},
+    )
+
+    assert response.status_code == 400, response.text
+    assert "Invalid budget_duration" in response.text
+    mock_prisma_client.db.litellm_endusertable.create.assert_not_awaited()
+
+
+def test_customer_new_accepts_a_normal_duration(mock_prisma_client, mock_user_api_key_auth):
+    mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW))
+    mock_prisma_client.db.litellm_budgettable.create = AsyncMock(
+        return_value=_row({"budget_id": "b1", "max_budget": 10.0})
+    )
+
+    response = client.post(
+        "/customer/new",
+        json={"user_id": "c1", "max_budget": 10.0, "budget_duration": "30d"},
+        headers={"Authorization": "Bearer k"},
+    )
+
+    assert response.status_code == 200, response.text
+
+
 def test_char_update_body(mock_prisma_client, mock_user_api_key_auth):
     mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(
         return_value=_row({"user_id": "c1", "blocked": False})
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 056c2d3657a..cd5a5d42b09 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
@@ -788,6 +788,68 @@ def test_update_internal_user_params_reset_spend_and_max_budget():
     assert "budget_duration" not in non_default_values  # Should not add default values
 
 
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+def test_update_internal_user_params_rejects_a_duration_that_never_advances(bad_duration):
+    """A zero-length window resets to "now", so the user row is due again the
+    moment it is written and the reset job re-reads it on every tick. Enough of
+    them fill each batch and starve other tenants' resets.
+    """
+    from fastapi import HTTPException
+
+    from litellm.proxy._types import UpdateUserRequest
+    from litellm.proxy.management_endpoints.internal_user_endpoints import (
+        _update_internal_user_params,
+    )
+
+    data = UpdateUserRequest(user_id="test_user_id", budget_duration=bad_duration)
+
+    with pytest.raises(HTTPException) as exc_info:
+        _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data)
+
+    assert exc_info.value.status_code == 400
+    assert "Invalid budget_duration" in str(exc_info.value.detail)
+
+
+def test_update_internal_user_params_accepts_a_normal_duration():
+    from litellm.proxy._types import UpdateUserRequest
+    from litellm.proxy.management_endpoints.internal_user_endpoints import (
+        _update_internal_user_params,
+    )
+
+    data = UpdateUserRequest(user_id="test_user_id", budget_duration="30d")
+
+    non_default_values = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data)
+
+    assert non_default_values["budget_duration"] == "30d"
+    assert non_default_values["budget_reset_at"] is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+async def test_new_user_rejects_a_duration_that_never_advances(mocker, bad_duration):
+    """/user/new must reject the same never-advancing durations /user/update does."""
+    from unittest.mock import AsyncMock, MagicMock
+
+    from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
+
+    mocker.patch("litellm.proxy.proxy_server.prisma_client", MagicMock())
+    duplicate_check = mocker.patch(
+        "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id",
+        new=AsyncMock(),
+    )
+    admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+
+    with pytest.raises(ProxyException) as exc_info:
+        await new_user(
+            data=NewUserRequest(budget_duration=bad_duration),
+            user_api_key_dict=admin,
+        )
+
+    assert str(exc_info.value.code) == "400"
+    assert "Invalid budget_duration" in str(exc_info.value.message)
+    duplicate_check.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_new_user_license_over_limit(mocker):
     """
diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
index 3ce25f7a934..8f151ed882c 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
@@ -2527,6 +2527,72 @@ def _setup_update_key_mocks(monkeypatch, mock_prisma_client):
     monkeypatch.setattr("litellm.store_audit_logs", False)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+async def test_update_key_rejects_a_duration_that_never_advances(monkeypatch, bad_duration):
+    """A zero-length window resets to "now", so the key row is due again the
+    moment it is written. The reset job re-reads such rows on every tick, and a
+    tenant with enough of them fills each batch and starves other tenants.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        update_key_fn,
+    )
+
+    hashed_token = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b"
+    key_in_db = LiteLLM_VerificationToken(token=hashed_token, user_id="test-user")
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=key_in_db
+    )
+    mock_prisma_client.update_data = AsyncMock()
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+
+    with pytest.raises(ProxyException) as exc_info:
+        await update_key_fn(
+            request=MagicMock(),
+            data=UpdateKeyRequest(key=hashed_token, budget_duration=bad_duration),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+            ),
+            litellm_changed_by=None,
+        )
+
+    assert str(exc_info.value.code) == "400"
+    assert "Invalid budget_duration" in str(exc_info.value.message)
+    mock_prisma_client.update_data.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+async def test_generate_key_rejects_a_duration_that_never_advances(monkeypatch, bad_duration):
+    """/key/generate must reject the same never-advancing durations /key/update does."""
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        generate_key_fn,
+    )
+
+    mock_prisma_client = AsyncMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+
+    with patch(
+        "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
+        new=AsyncMock(),
+    ) as mock_generate:
+        with pytest.raises(ProxyException) as exc_info:
+            await generate_key_fn(
+                data=GenerateKeyRequest(budget_duration=bad_duration),
+                user_api_key_dict=UserAPIKeyAuth(
+                    user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234"
+                ),
+            )
+
+    assert str(exc_info.value.code) == "400"
+    assert "Invalid budget_duration" in str(exc_info.value.message)
+    mock_generate.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_update_key_by_alias_only(monkeypatch):
     """
@@ -8081,7 +8147,7 @@ async def test_key_with_budget_id_does_not_store_budget_duration():
     budget_duration, the key does NOT get budget_duration stored on it.
 
     Keys with budget_id follow their linked budget tier's reset schedule;
-    reset_budget_for_keys_linked_to_budgets() resets them when the tier resets.
+    reset_budget_for_litellm_budget_table() resets them when the tier resets.
     This avoids duplicating budget_duration on keys so tier updates apply
     automatically to all linked keys.
     """
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index f17a6dbd380..6abc40eb28e 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -380,6 +380,66 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth):
         app.dependency_overrides = {}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field", ["budget_duration", "team_member_budget_duration"])
+@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
+async def test_new_team_rejects_a_duration_that_never_advances(
+    mock_db_client, mock_admin_auth, field, bad_duration
+):
+    """A zero-length window resets to "now", so the team row is due again the
+    moment it is written. The reset job re-reads such rows on every tick, and a
+    tenant with enough of them fills each batch and starves other tenants.
+    """
+    from fastapi import Request
+
+    from litellm.proxy._types import NewTeamRequest
+    from litellm.proxy.management_endpoints.team_endpoints import new_team
+
+    mock_db_client.db = MagicMock()
+    mock_team_create = AsyncMock()
+    mock_db_client.db.litellm_teamtable = MagicMock()
+    mock_db_client.db.litellm_teamtable.create = mock_team_create
+
+    with pytest.raises(ProxyException) as exc_info:
+        await new_team(
+            data=NewTeamRequest(team_alias="my-team", **{field: bad_duration}),
+            http_request=MagicMock(spec=Request),
+            user_api_key_dict=mock_admin_auth,
+        )
+
+    assert str(exc_info.value.code) == "400"
+    assert "Invalid budget_duration" in str(exc_info.value.message)
+    mock_team_create.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("field", ["budget_duration", "team_member_budget_duration"])
+async def test_update_team_rejects_a_duration_that_never_advances(
+    mock_db_client, mock_admin_auth, field
+):
+    """/team/update must reject the same never-advancing durations /team/new does."""
+    from fastapi import Request
+
+    from litellm.proxy._types import UpdateTeamRequest
+    from litellm.proxy.management_endpoints.team_endpoints import update_team
+
+    mock_db_client.db = MagicMock()
+    mock_find_unique = AsyncMock(return_value=None)
+    mock_db_client.db.litellm_teamtable = MagicMock()
+    mock_db_client.db.litellm_teamtable.find_unique = mock_find_unique
+
+    with pytest.raises(ProxyException) as exc_info:
+        await update_team(
+            data=UpdateTeamRequest(team_id="team-1", **{field: "0s"}),
+            http_request=MagicMock(spec=Request),
+            user_api_key_dict=mock_admin_auth,
+        )
+
+    assert str(exc_info.value.code) == "400"
+    assert "Invalid budget_duration" in str(exc_info.value.message)
+    mock_find_unique.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth):
     """
diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py
index 35f102bbb9d..c270a570ad9 100644
--- a/tests/test_litellm/repositories/test_unit_of_work.py
+++ b/tests/test_litellm/repositories/test_unit_of_work.py
@@ -3,7 +3,10 @@ from typing import Any, Dict, List, Mapping, Tuple
 
 import pytest
 
-from litellm.repositories.unit_of_work import spend_reset_unit_of_work
+from litellm.repositories.unit_of_work import (
+    budget_cascade_unit_of_work,
+    spend_reset_unit_of_work,
+)
 
 
 class FakeBatchTable:
@@ -14,6 +17,9 @@ class FakeBatchTable:
     def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> None:
         self._calls.append((self._table_name, dict(where), dict(data)))
 
+    def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> None:
+        self._calls.append((f"{self._table_name}.update_many", dict(where), dict(data)))
+
 
 class FakeBatch:
     def __init__(self):
@@ -22,6 +28,11 @@ class FakeBatch:
         self.litellm_verificationtoken = FakeBatchTable("litellm_verificationtoken", self.calls)
         self.litellm_usertable = FakeBatchTable("litellm_usertable", self.calls)
         self.litellm_teamtable = FakeBatchTable("litellm_teamtable", self.calls)
+        self.litellm_budgettable = FakeBatchTable("litellm_budgettable", self.calls)
+        self.litellm_teammembership = FakeBatchTable("litellm_teammembership", self.calls)
+        self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls)
+        self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls)
+        self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls)
 
     async def commit(self) -> None:
         self.commit_count += 1
@@ -64,3 +75,53 @@ async def test_empty_block_still_commits_the_batch():
 
     assert batch.commit_count == 1
     assert batch.calls == []
+
+
+async def test_budget_cascade_dependents_and_window_advance_share_one_batch():
+    batch = FakeBatch()
+    reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc)
+    linked = {"budget_id": {"in": ["budget-1"]}}
+
+    async with budget_cascade_unit_of_work(lambda: batch) as uow:
+        uow.team_memberships.queue_spend_zero(where=linked)
+        uow.keys.queue_spend_zero(where=linked)
+        uow.organizations.queue_spend_zero(where=linked)
+        uow.tags.queue_spend_zero(where=linked)
+        uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}})
+        uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at)
+        assert batch.commit_count == 0
+
+    assert batch.commit_count == 1
+    assert batch.calls == [
+        ("litellm_teammembership.update_many", linked, {"spend": 0}),
+        ("litellm_verificationtoken.update_many", linked, {"spend": 0}),
+        ("litellm_organizationtable.update_many", linked, {"spend": 0}),
+        ("litellm_tagtable.update_many", linked, {"spend": 0}),
+        ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}),
+        ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}),
+    ]
+
+
+async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk():
+    """A tier deleted between the read and the commit must not abort the batch:
+    ``update`` raises P2025 on a missing row and takes every other write in the
+    chunk down with it, while ``update_many`` just matches nothing."""
+    batch = FakeBatch()
+
+    async with budget_cascade_unit_of_work(lambda: batch) as uow:
+        uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=datetime.now(timezone.utc))
+
+    assert [call[0] for call in batch.calls] == ["litellm_budgettable.update_many"]
+
+
+async def test_budget_cascade_raising_inside_block_skips_commit():
+    """A failure part-way through must leave budget_reset_at where it was, so
+    the tier is still due on the next tick."""
+    batch = FakeBatch()
+
+    with pytest.raises(RuntimeError, match="boom"):
+        async with budget_cascade_unit_of_work(lambda: batch) as uow:
+            uow.team_memberships.queue_spend_zero(where={"budget_id": {"in": ["budget-1"]}})
+            raise RuntimeError("boom")
+
+    assert batch.commit_count == 0
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 585f8cd77f9..c990ae52ff2 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -1,9 +1,9 @@
 {
   "LIT001": {
-    "limit": 23064
+    "limit": 23057
   },
   "LIT002": {
-    "limit": 27166
+    "limit": 27156
   },
   "LIT003": {
     "limit": 269
@@ -27,9 +27,9 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16753
+    "limit": 16744
   },
   "LIT011": {
-    "limit": 5598
+    "limit": 5596
   }
 }

From 76ad1c319de66c9213fa5295677b1baadd01717f Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Mon, 10 Aug 2026 15:09:59 -0700
Subject: [PATCH 6/6] feat(proxy): add GET /v1/indexes to list vector store
 indexes (#36289)

* fix(scripts): stop type-discipline checker reading Literal strings as forward refs

The checker re-parsed every string constant inside an annotation as a
forward reference, so Literal["list"] was counted as the mutable list
type. Skip Literal subtrees and ratchet the LIT001 ceiling down to the
corrected count.

* fix(proxy): keep lazy openapi snapshot fragments for transitively imported features

generate_snapshot skipped register_fn for any feature module already in
sys.modules, so a module pulled in transitively by an earlier feature
never mounted its routes and its fragment silently vanished on regen
(vector_store_management). Route collection also matched path_prefixes
only, dropping suffix-matched routes from fragments. Register every
feature and collect routes with feat.matches, mirroring the runtime
loader.

* feat(proxy): add GET /v1/indexes to list vector store indexes

/v1/indexes was POST-only, so indexes created through it could never be
viewed again. Add an admin-only list endpoint returning the stored index
rows newest first, fix the stale index_create docstring curl, and
regenerate the lazy openapi snapshot and dashboard schema types.

* chore(proxy): defer lazy openapi snapshot catch-up regen to a follow-up

Reverts _lazy_openapi_snapshot.json and schema.d.ts to the staging
versions. The snapshot was months stale, so regenerating it here buried
the actual change under ten thousand generated lines. A follow-up will
land the regen together with CI enforcement that keeps the snapshot
current. Until then GET /v1/indexes is served but absent from the
dashboard's generated types, which the UI step needs anyway.

* fix(proxy): use Annotated dependency to avoid new B008 violation
---
 litellm/proxy/_lazy_openapi_snapshot.py       | 11 ++-
 .../proxy/vector_store_endpoints/endpoints.py | 49 ++++++++--
 litellm/proxy/vector_store_endpoints/utils.py |  2 +-
 litellm/types/vector_stores.py                |  5 +
 scripts/check_type_discipline.py              | 42 ++++++---
 .../proxy/test_lazy_openapi_snapshot.py       | 76 ++++++++++-----
 .../test_vector_store_endpoints.py            | 92 ++++++++++++++++++-
 .../test_check_type_discipline.py             | 11 +++
 8 files changed, 235 insertions(+), 53 deletions(-)

diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py
index 36ffc819774..41359d44b27 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.py
+++ b/litellm/proxy/_lazy_openapi_snapshot.py
@@ -3,8 +3,11 @@ Per-feature OpenAPI snapshot for lazy-loaded routers.
 
 The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot`
 and consumed at runtime so /openapi.json can show full route info for unloaded
-features without importing them. CI verifies the file is current and surfaces
-any drift as a neutral check.
+features without importing them. No CI job regenerates this file; drift surfaces
+only indirectly through check-ui-api-types.yml, which rebuilds schema.d.ts from
+app.openapi() with the committed snapshot injected. After changing any lazily
+loaded route or this generator, rerun the module and commit the JSON, then run
+`npm run gen:api` in ui/litellm-dashboard and commit schema.d.ts.
 """
 
 import json
@@ -89,8 +92,6 @@ def generate_snapshot() -> dict[str, dict]:
     from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids
 
     for feat in LAZY_FEATURES:
-        if feat.module_path in sys.modules:
-            continue
         try:
             module = importlib.import_module(feat.module_path)
             feat.register_fn(app, module)
@@ -100,7 +101,7 @@ def generate_snapshot() -> dict[str, dict]:
     fragments: Final[dict[str, dict]] = {}
     used_operation_ids: Final[set[str]] = set()
     for feat in LAZY_FEATURES:
-        feat_routes = [r for r in app.routes if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)]
+        feat_routes = [r for r in app.routes if feat.matches(getattr(r, "path", ""))]
         if not feat_routes:
             continue
         _stabilize_multi_method_route_ids(feat_routes)
diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py
index c2483c81d6c..b497247f576 100644
--- a/litellm/proxy/vector_store_endpoints/endpoints.py
+++ b/litellm/proxy/vector_store_endpoints/endpoints.py
@@ -1,4 +1,4 @@
-from typing import Any, Final
+from typing import Annotated, Any, Final
 
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
 
@@ -18,7 +18,8 @@ from litellm.proxy.vector_store_endpoints.utils import (
     get_litellm_managed_vector_store,
 )
 from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository
-from litellm.types.vector_stores import IndexCreateRequest
+from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
+from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry
 
 router: Final = APIRouter()
 ########################################################
@@ -549,14 +550,15 @@ async def index_create(
     Create an index. Just writes the index to the database.
 
     ```bash
-    curl -L -X POST 'http://0.0.0.0:4000/indexes/create' \
+    curl -L -X POST 'http://0.0.0.0:4000/v1/indexes' \
         -H 'Content-Type: application/json' \
         -H 'Authorization: Bearer sk-1234' \
-        -H 'LiteLLM-Beta: indexes_beta=v1' \
-        -d '{ 
+        -d '{
             "index_name": "dall-e-3",
-            "vector_store_index": "real-index-name",
-            "vector_store_name": "azure-ai-search"
+            "litellm_params": {
+                "vector_store_index": "real-index-name",
+                "vector_store_name": "azure-ai-search"
+            }
         }'
     ```
     """
@@ -592,3 +594,36 @@ async def index_create(
     new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(data=jsonify_object(index_data))
 
     return new_index.model_dump()
+
+
+@router.get(
+    "/v1/indexes",
+    dependencies=[Depends(user_api_key_auth)],
+    response_model=IndexListResponse,
+)
+async def index_list(
+    user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+) -> IndexListResponse:
+    """
+    List all vector store indexes. Proxy admin only.
+
+    ```bash
+    curl -L -X GET 'http://0.0.0.0:4000/v1/indexes' \
+        -H 'Authorization: Bearer sk-1234'
+    ```
+    """
+    from litellm.proxy.proxy_server import prisma_client
+
+    assert_proxy_admin_for_vector_store_index_management(
+        user_api_key_dict,
+        operation="list",
+    )
+
+    if prisma_client is None:
+        raise HTTPException(
+            status_code=500,
+            detail=CommonProxyErrors.db_not_connected_error.value,
+        )
+
+    indexes: Final = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db(prisma_client)
+    return IndexListResponse(data=indexes)
diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py
index 402fba65558..94ba7c06cad 100644
--- a/litellm/proxy/vector_store_endpoints/utils.py
+++ b/litellm/proxy/vector_store_endpoints/utils.py
@@ -41,7 +41,7 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
 def assert_proxy_admin_for_vector_store_index_management(
     user_api_key_dict: UserAPIKeyAuth,
     *,
-    operation: Literal["create", "delete", "update"] = "create",
+    operation: Literal["create", "delete", "update", "list"] = "create",
 ) -> None:
     """Raise 403 unless the caller is a proxy admin."""
     if _is_proxy_admin(user_api_key_dict):
diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py
index d1d4a39da1e..474c652ff3a 100644
--- a/litellm/types/vector_stores.py
+++ b/litellm/types/vector_stores.py
@@ -277,6 +277,11 @@ class LiteLLM_ManagedVectorStoreIndex(BaseModel):
     updated_by: str | None = None
 
 
+class IndexListResponse(BaseModel):
+    object: Literal["list"] = "list"
+    data: tuple[LiteLLM_ManagedVectorStoreIndex, ...]
+
+
 class VectorStoreIndexType(str, Enum):
     """Type of vector store index"""
 
diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py
index 65d0424fb5a..92eb7ef55a3 100644
--- a/scripts/check_type_discipline.py
+++ b/scripts/check_type_discipline.py
@@ -252,26 +252,40 @@ def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, .
 # --------------------------------------------------------------------------- #
  
  
-def mutable_names_in(annotation: ast.expr) -> Iterator[str]:
+def _is_literal_subscript(node: ast.AST) -> bool:
+    if not isinstance(node, ast.Subscript):
+        return False
+    base: Final = node.value
+    return (isinstance(base, ast.Name) and base.id == "Literal") or (
+        isinstance(base, ast.Attribute) and base.attr == "Literal"
+    )
+
+
+def mutable_names_in(annotation: ast.AST) -> Iterator[str]:
     """Yield mutable-collection names anywhere inside an annotation expression.
 
     Matches bare names (`dict`, `MutableMapping`) and dotted access (`typing.Dict`,
     `collections.deque`, `collections.abc.MutableMapping`), descends through nesting
     (`Mapping[str, list[int]]`, `tuple[set[int], ...]`) and string forward references.
+    Skips `Literal[...]` subtrees: their string arguments are values, not forward
+    references, so `Literal["list"]` is not the `list` type.
     """
-    for node in ast.walk(annotation):
-        if isinstance(node, ast.Name) and node.id in MUTABLE_COLLECTIONS:
-            yield node.id
-        elif isinstance(node, ast.Attribute) and node.attr in MUTABLE_COLLECTIONS:
-            yield node.attr
-        elif isinstance(node, ast.Constant):
-            value: object = node.value  # forward references arrive as string constants
-            if isinstance(value, str):
-                try:
-                    inner = ast.parse(value, mode="eval").body
-                except SyntaxError:
-                    continue
-                yield from mutable_names_in(inner)
+    if _is_literal_subscript(annotation):
+        return
+    if isinstance(annotation, ast.Name) and annotation.id in MUTABLE_COLLECTIONS:
+        yield annotation.id
+    elif isinstance(annotation, ast.Attribute) and annotation.attr in MUTABLE_COLLECTIONS:
+        yield annotation.attr
+    elif isinstance(annotation, ast.Constant):
+        value: object = annotation.value  # forward references arrive as string constants
+        if isinstance(value, str):
+            try:
+                inner = ast.parse(value, mode="eval").body
+            except SyntaxError:
+                return
+            yield from mutable_names_in(inner)
+    for child in ast.iter_child_nodes(annotation):
+        yield from mutable_names_in(child)
  
  
 def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation:
diff --git a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py
index 64cb931888b..79330b0e3a6 100644
--- a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py
+++ b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py
@@ -1,6 +1,7 @@
 import sys
 from types import ModuleType, SimpleNamespace
 
+from litellm.proxy._lazy_features import LazyFeature
 from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids
 
 
@@ -22,22 +23,20 @@ def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch):
 
     fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features")
     fake_lazy_features_module.LAZY_FEATURES = [
-        SimpleNamespace(
+        LazyFeature(
             name="feature-a",
             module_path="fake_feature_a",
             path_prefixes=("/feature-a",),
             register_fn=lambda app, module: None,
         ),
-        SimpleNamespace(
+        LazyFeature(
             name="feature-b",
             module_path="fake_feature_b",
             path_prefixes=("/feature-b",),
             register_fn=lambda app, module: None,
         ),
     ]
-    monkeypatch.setitem(
-        sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module
-    )
+    monkeypatch.setitem(sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module)
 
     def fake_get_openapi(title, version, routes):
         path = routes[0].path
@@ -58,30 +57,59 @@ def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch):
 
     fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server")
     fake_proxy_server_module.app = fake_app
-    fake_proxy_server_module.ensure_unique_openapi_operation_ids = (
-        fake_ensure_unique_openapi_operation_ids
-    )
-    monkeypatch.setitem(
-        sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module
-    )
+    fake_proxy_server_module.ensure_unique_openapi_operation_ids = fake_ensure_unique_openapi_operation_ids
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module)
     monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi)
 
     fragments = _lazy_openapi_snapshot.generate_snapshot()
 
-    assert (
-        fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"]
-        == "shared_operation_id_get"
-    )
-    assert (
-        fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"]
-        == "shared_operation_id_get_2"
-    )
-    assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["tags"] == [
-        "feature-a"
-    ]
-    assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["tags"] == [
-        "feature-b"
+    assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"] == "shared_operation_id_get"
+    assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"] == "shared_operation_id_get_2"
+    assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["tags"] == ["feature-a"]
+    assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["tags"] == ["feature-b"]
+
+
+def test_generate_snapshot_registers_transitively_imported_modules(monkeypatch):
+    """A feature module already in sys.modules (pulled in transitively by an
+    earlier feature) must still get register_fn called, else its routes never
+    mount and its fragment silently vanishes from the snapshot. Fragment
+    collection must also honor path_suffixes, not just prefixes."""
+    from litellm.proxy import _lazy_openapi_snapshot
+
+    fake_app = SimpleNamespace(title="LiteLLM test", version="0.0.0", routes=[])
+
+    fake_module = ModuleType("fake_transitive_feature")
+    monkeypatch.setitem(sys.modules, "fake_transitive_feature", fake_module)
+
+    def register_fn(app, module):
+        app.routes.append(SimpleNamespace(path="/transitive/items"))
+        app.routes.append(SimpleNamespace(path="/v1/{param}/deep/leaf"))
+
+    fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features")
+    fake_lazy_features_module.LAZY_FEATURES = [
+        LazyFeature(
+            name="transitive",
+            module_path="fake_transitive_feature",
+            path_prefixes=("/transitive",),
+            path_suffixes=("/deep/leaf",),
+            register_fn=register_fn,
+        )
     ]
+    monkeypatch.setitem(sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module)
+
+    def fake_get_openapi(title, version, routes):
+        return {"paths": {route.path: {"get": {"operationId": f"op{i}_get"}} for i, route in enumerate(routes)}}
+
+    fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server")
+    fake_proxy_server_module.app = fake_app
+    fake_proxy_server_module.ensure_unique_openapi_operation_ids = lambda schema, reserved_operation_ids: schema
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module)
+    monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi)
+
+    fragments = _lazy_openapi_snapshot.generate_snapshot()
+
+    assert fragments["transitive"]["paths"]["/transitive/items"]["get"]["tags"] == ["transitive"]
+    assert "/v1/{param}/deep/leaf" in fragments["transitive"]["paths"]
 
 
 def test_normalize_operation_ids_uses_each_http_method():
diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
index e7de8b54e4e..02ca64e5fb8 100644
--- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
+++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py
@@ -16,10 +16,11 @@ import litellm
 from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
     LiteLLM_ManagedVectorStore,
 )
-from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.vector_store_endpoints.endpoints import (
     _update_request_data_with_litellm_managed_vector_store_registry,
     index_create,
+    index_list,
 )
 from litellm.proxy.vector_store_files_endpoints.endpoints import (
     _update_request_data_with_model_routing_hint,
@@ -37,7 +38,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
     is_allowed_to_call_vector_store_endpoint,
     is_allowed_to_call_vector_store_files_endpoint,
 )
-from litellm.types.vector_stores import IndexCreateRequest
+from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
 from litellm.types.utils import LlmProviders
 
 
@@ -1316,6 +1317,93 @@ class TestIndexCreate:
         mock_prisma.db.litellm_managedvectorstoreindextable.create.assert_awaited_once()
 
 
+class TestIndexList:
+    def _admin(self) -> UserAPIKeyAuth:
+        return UserAPIKeyAuth(
+            token="sk-test",
+            key_name="sk-...test",
+            user_role=LitellmUserRoles.PROXY_ADMIN,
+            user_id="admin-user",
+        )
+
+    def _index_row(self, index_id: str, index_name: str) -> dict:
+        return {
+            "id": index_id,
+            "index_name": index_name,
+            "litellm_params": {
+                "vector_store_index": f"real-{index_name}",
+                "vector_store_name": "azure-ai-search",
+            },
+            "index_info": None,
+            "created_at": datetime(2026, 1, 2, tzinfo=timezone.utc),
+            "created_by": "admin-user",
+            "updated_at": datetime(2026, 1, 2, tzinfo=timezone.utc),
+            "updated_by": "admin-user",
+        }
+
+    @pytest.mark.asyncio
+    async def test_index_list_requires_admin(self):
+        """Index topology must never reach non-admins, not even via a DB read."""
+        mock_prisma = MagicMock()
+        mock_prisma.db.litellm_managedvectorstoreindextable.find_many = AsyncMock()
+
+        with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
+            with pytest.raises(HTTPException) as exc_info:
+                await index_list(
+                    user_api_key_dict=UserAPIKeyAuth(
+                        token="sk-test",
+                        key_name="sk-...test",
+                        user_role=LitellmUserRoles.INTERNAL_USER,
+                    )
+                )
+
+        assert exc_info.value.status_code == 403
+        assert "Only proxy admins can list" in exc_info.value.detail
+        mock_prisma.db.litellm_managedvectorstoreindextable.find_many.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_index_list_requires_db_connection(self):
+        with patch("litellm.proxy.proxy_server.prisma_client", None):
+            with pytest.raises(HTTPException) as exc_info:
+                await index_list(user_api_key_dict=self._admin())
+
+        assert exc_info.value.status_code == 500
+        assert CommonProxyErrors.db_not_connected_error.value in exc_info.value.detail
+
+    @pytest.mark.asyncio
+    async def test_index_list_returns_db_rows_newest_first(self):
+        """Rows round-trip into typed models and DB ordering (created_at desc) is requested."""
+        rows = [
+            self._index_row("idx-2", "index-b"),
+            self._index_row("idx-1", "index-a"),
+        ]
+        mock_prisma = MagicMock()
+        mock_prisma.db.litellm_managedvectorstoreindextable.find_many = AsyncMock(return_value=rows)
+
+        with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
+            result = await index_list(user_api_key_dict=self._admin())
+
+        assert isinstance(result, IndexListResponse)
+        assert result.object == "list"
+        assert [index.index_name for index in result.data] == ["index-b", "index-a"]
+        assert result.data[0].litellm_params.vector_store_index == "real-index-b"
+        assert result.data[0].litellm_params.vector_store_name == "azure-ai-search"
+        assert result.data[1].litellm_params.vector_store_index == "real-index-a"
+        mock_prisma.db.litellm_managedvectorstoreindextable.find_many.assert_awaited_once_with(
+            order={"created_at": "desc"}
+        )
+
+    def test_get_v1_indexes_route_registered(self):
+        from litellm.proxy.vector_store_endpoints.endpoints import router
+
+        routes = [
+            (method, getattr(route, "path", None))
+            for route in router.routes
+            for method in (getattr(route, "methods", None) or ())
+        ]
+        assert ("GET", "/v1/indexes") in routes
+
+
 class TestIsAllowedToCallVectorStoreFilesEndpoint:
     def _mock_provider_config(self):
         provider_config = MagicMock()
diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py
index 4b8533df604..25131088e9a 100644
--- a/tests/test_litellm/test_check_type_discipline.py
+++ b/tests/test_litellm/test_check_type_discipline.py
@@ -117,6 +117,17 @@ def test_typing_alias_and_forward_ref_annotations_are_flagged(tmp_path):
     assert "LIT001" in _codes(tmp_path, 'x: "dict[str, int]"\n')
 
 
+def test_literal_string_args_are_values_not_forward_refs(tmp_path):
+    assert "LIT001" not in _codes(tmp_path, 'from typing import Literal\nx: Literal["list"] = "list"\n')
+    assert "LIT001" not in _codes(
+        tmp_path,
+        'from typing import Literal\ndef f(op: Literal["create", "list"] = "create") -> None:\n    return None\n',
+    )
+    assert "LIT001" not in _codes(tmp_path, 'import typing\nx: typing.Literal["dict"] = "dict"\n')
+    assert "LIT001" in _codes(tmp_path, 'from typing import Literal\nx: dict[str, Literal["a"]]\n')
+    assert "LIT001" in _codes(tmp_path, "x: \"Literal['x'] | list[int]\"\n")
+
+
 def test_readonly_annotations_are_clean(tmp_path):
     for ann in ("Mapping[str, int]", "Sequence[int]", "tuple[int, ...]", "frozenset[int]"):
         assert "LIT001" not in _codes(tmp_path, f"from typing import Mapping, Sequence\nx: {ann}\n")