From e5e71f55aa0783a1db723926939591bce7c03654 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:23:16 -0700 Subject: [PATCH 1/6] fix(cost-map): sync openrouter deepseek v4 and glm-5.3-flash prices, add mistral-large-2512 (#43090) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 38 ++++++++++++++----- model_prices_and_context_window.json | 38 ++++++++++++++----- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3e2aa011164..f12c8563392 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41395,20 +41395,20 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.1263e-07, + "input_cost_per_token": 8.44944e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.82526e-06, + "output_cost_per_token": 1.689888e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.60525e-08, + "cache_read_input_token_cost": 7.0412e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65906,9 +65906,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.5e-08, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2.85e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -66657,9 +66657,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.4e-08, - "output_cost_per_token": 1.68e-07, - "cache_read_input_token_cost": 1.68e-08, + "input_cost_per_token": 4.9e-08, + "output_cost_per_token": 9.8e-08, + "cache_read_input_token_cost": 9.8e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -73050,6 +73050,26 @@ "supports_vision": true, "supports_web_search": false }, + "openrouter/mistralai/mistral-large-2512": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 209715, + "max_tokens": 209715, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, "openrouter/mistralai/mistral-large-2512:batch": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3e2aa011164..f12c8563392 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41395,20 +41395,20 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.1263e-07, + "input_cost_per_token": 8.44944e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.82526e-06, + "output_cost_per_token": 1.689888e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.60525e-08, + "cache_read_input_token_cost": 7.0412e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65906,9 +65906,9 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 4.5e-08, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 2.85e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -66657,9 +66657,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.4e-08, - "output_cost_per_token": 1.68e-07, - "cache_read_input_token_cost": 1.68e-08, + "input_cost_per_token": 4.9e-08, + "output_cost_per_token": 9.8e-08, + "cache_read_input_token_cost": 9.8e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -73050,6 +73050,26 @@ "supports_vision": true, "supports_web_search": false }, + "openrouter/mistralai/mistral-large-2512": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 5e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 262144, + "max_output_tokens": 209715, + "max_tokens": 209715, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": false, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": false + }, "openrouter/mistralai/mistral-large-2512:batch": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, From cbe342171cb30e98bcfe6d555ba60489eff46b9d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:27:00 -0700 Subject: [PATCH 2/6] feat(cost-map): add fireworks glm-5p3 US-only rows and kimi-k3-us priority prices (#43092) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 66 +++++++++++++++++++ model_prices_and_context_window.json | 66 +++++++++++++++++++ 2 files changed, 132 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f12c8563392..c58ea51db9f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -60319,13 +60319,16 @@ }, "fireworks_ai/kimi-k3-us": { "cache_read_input_token_cost": 4.5e-07, + "cache_read_input_token_cost_priority": 5.625e-07, "input_cost_per_token": 4.5e-06, + "input_cost_per_token_priority": 5.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.25e-05, + "output_cost_per_token_priority": 2.8125e-05, "reasoning_effort_levels": [ "low", "high", @@ -60527,13 +60530,16 @@ }, "fireworks_ai/accounts/fireworks/routers/kimi-k3-us": { "cache_read_input_token_cost": 4.5e-07, + "cache_read_input_token_cost_priority": 5.625e-07, "input_cost_per_token": 4.5e-06, + "input_cost_per_token_priority": 5.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.25e-05, + "output_cost_per_token_priority": 2.8125e-05, "reasoning_effort_levels": [ "low", "high", @@ -63292,6 +63298,22 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/glm-5p3": { "cache_read_input_token_cost": 2.6e-07, "cache_read_input_token_cost_priority": 3.25e-07, @@ -63311,6 +63333,22 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/glm-5p3-us": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { "cache_read_input_token_cost": 3.9e-07, "input_cost_per_token": 2.1e-06, @@ -63358,6 +63396,20 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": { + "cache_read_input_token_cost": 4.5e-08, + "input_cost_per_token": 2.25e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_priority": 3.75e-08, @@ -63375,6 +63427,20 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/glm-5p3-flash-us": { + "cache_read_input_token_cost": 4.5e-08, + "input_cost_per_token": 2.25e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/inkling": { "cache_read_input_token_cost": 1.7e-07, "input_cost_per_token": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f12c8563392..c58ea51db9f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -60319,13 +60319,16 @@ }, "fireworks_ai/kimi-k3-us": { "cache_read_input_token_cost": 4.5e-07, + "cache_read_input_token_cost_priority": 5.625e-07, "input_cost_per_token": 4.5e-06, + "input_cost_per_token_priority": 5.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.25e-05, + "output_cost_per_token_priority": 2.8125e-05, "reasoning_effort_levels": [ "low", "high", @@ -60527,13 +60530,16 @@ }, "fireworks_ai/accounts/fireworks/routers/kimi-k3-us": { "cache_read_input_token_cost": 4.5e-07, + "cache_read_input_token_cost_priority": 5.625e-07, "input_cost_per_token": 4.5e-06, + "input_cost_per_token_priority": 5.625e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2.25e-05, + "output_cost_per_token_priority": 2.8125e-05, "reasoning_effort_levels": [ "low", "high", @@ -63292,6 +63298,22 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-us": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/glm-5p3": { "cache_read_input_token_cost": 2.6e-07, "cache_read_input_token_cost_priority": 3.25e-07, @@ -63311,6 +63333,22 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/glm-5p3-us": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { "cache_read_input_token_cost": 3.9e-07, "input_cost_per_token": 2.1e-06, @@ -63358,6 +63396,20 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-flash-us": { + "cache_read_input_token_cost": 4.5e-08, + "input_cost_per_token": 2.25e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_priority": 3.75e-08, @@ -63375,6 +63427,20 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/glm-5p3-flash-us": { + "cache_read_input_token_cost": 4.5e-08, + "input_cost_per_token": 2.25e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/inkling": { "cache_read_input_token_cost": 1.7e-07, "input_cost_per_token": 1e-06, From 165dbc22437915993d45cd96252c1fea1f969fbc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:28:48 -0700 Subject: [PATCH 3/6] feat(ui): offer reset of custom member budgets when team default changes (#42835) --- .../team/ResetMemberBudgetsDialog.test.tsx | 103 +++ .../team/ResetMemberBudgetsDialog.tsx | 71 ++ .../team/TeamInfo.integration.test.tsx | 722 ++++++++++++++++++ .../src/components/team/TeamInfo.test.tsx | 6 + .../src/components/team/TeamInfo.tsx | 69 +- .../components/team/TeamMemberTab.test.tsx | 38 + .../src/components/team/TeamMemberTab.tsx | 5 +- .../components/team/memberBudgetReset.test.ts | 101 +++ .../src/components/team/memberBudgetReset.ts | 38 + .../team/useMemberBudgetReset.test.tsx | 345 +++++++++ .../components/team/useMemberBudgetReset.ts | 134 ++++ 11 files changed, 1623 insertions(+), 9 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx create mode 100644 ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx create mode 100644 ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts create mode 100644 ui/litellm-dashboard/src/components/team/memberBudgetReset.ts create mode 100644 ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx create mode 100644 ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts diff --git a/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx new file mode 100644 index 00000000000..7132ec13fbb --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.test.tsx @@ -0,0 +1,103 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog"; +import type { MemberBudgetResetState } from "./useMemberBudgetReset"; + +const pendingFor = (memberCount: number, newBudget = 10) => ({ + teamId: "team-123", + updateData: {}, + userIds: Array.from({ length: memberCount }, (_, i) => `u-${i}`), + newBudget, +}); + +const promptingWith = (memberCount: number, newBudget = 10): MemberBudgetResetState => ({ + phase: "prompting", + pending: pendingFor(memberCount, newBudget), +}); + +const defaultHandlers = () => ({ + onReset: vi.fn(), + onRetry: vi.fn(), + onKeep: vi.fn(), + onDismiss: vi.fn(), +}); + +describe("ResetMemberBudgetsDialog", () => { + it("names the members and the new default so the prompt is self-explanatory", () => { + render(); + + expect( + screen.getByText( + "3 members have a custom budget, so the new team default of $10.00 will not apply to them. " + + "Reset them to the default, or keep the custom budgets?", + ), + ).toBeInTheDocument(); + }); + + it("switches to singular for one member", () => { + render(); + + expect( + screen.getByText( + "1 member has a custom budget, so the new team default of $10.00 will not apply to that member. " + + "Reset it to the default, or keep the custom budget?", + ), + ).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Reset to $10.00" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Keep custom budget" })).toBeInTheDocument(); + }); + + it("formats large budgets like the rest of the team page", () => { + render(); + + expect(screen.getByRole("button", { name: "Reset all to $1,000.00" })).toBeInTheDocument(); + expect(screen.getByText(/new team default of \$1,000\.00/)).toBeInTheDocument(); + }); + + it("routes each choice to the matching action", async () => { + const handlers = defaultHandlers(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: "Reset all to $10.00" })); + expect(handlers.onReset).toHaveBeenCalledTimes(1); + + await user.click(screen.getByRole("button", { name: "Keep custom budgets" })); + expect(handlers.onKeep).toHaveBeenCalledTimes(1); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + expect(handlers.onDismiss).toHaveBeenCalledTimes(1); + }); + + it("locks every button while the reset is in flight", () => { + const state: MemberBudgetResetState = { phase: "resetting", pending: pendingFor(2), attempted: 0 }; + render(); + + expect(screen.getByRole("button", { name: "Reset all to $10.00" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Keep custom budgets" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled(); + }); + + it("offers retry or cancel after a failed reset", async () => { + const handlers = defaultHandlers(); + const user = userEvent.setup(); + const state: MemberBudgetResetState = { phase: "resetFailed", pending: pendingFor(2), attempted: 0 }; + render(); + + await user.click(screen.getByRole("button", { name: "Retry reset" })); + expect(handlers.onRetry).toHaveBeenCalledTimes(1); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + expect(handlers.onDismiss).toHaveBeenCalledTimes(1); + + expect(screen.queryByRole("button", { name: "Keep custom budgets" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Reset all to $10.00" })).not.toBeInTheDocument(); + }); + + it("renders nothing while idle", () => { + render(); + + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx new file mode 100644 index 00000000000..95714e3df33 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/ResetMemberBudgetsDialog.tsx @@ -0,0 +1,71 @@ +import React from "react"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { pluralize } from "./memberBudgetReset"; +import type { MemberBudgetResetState } from "./useMemberBudgetReset"; + +interface ResetMemberBudgetsDialogProps { + state: MemberBudgetResetState; + onReset: () => void; + onRetry: () => void; + onKeep: () => void; + onDismiss: () => void; +} + +export default function ResetMemberBudgetsDialog({ + state, + onReset, + onRetry, + onKeep, + onDismiss, +}: ResetMemberBudgetsDialogProps) { + const open = state.phase !== "idle"; + const busy = state.phase === "resetting"; + const failed = state.phase === "resetFailed"; + const memberCount = state.phase === "idle" ? 0 : state.pending.userIds.length; + const newBudget = state.phase === "idle" ? 0 : state.pending.newBudget; + + return ( + { + if (!nextOpen && !busy) onDismiss(); + }} + > + + + Reset member budgets? + +

+ {memberCount} {pluralize(memberCount, "member has", "members have")} a custom budget, so the new team default + of ${formatNumberWithCommas(newBudget, 2)} will not apply to {pluralize(memberCount, "that member", "them")}. + Reset {pluralize(memberCount, "it", "them")} to the default, or keep the custom{" "} + {pluralize(memberCount, "budget", "budgets")}? +

+ + {failed ? ( + <> + + + + ) : ( + <> + + + + + )} + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx new file mode 100644 index 00000000000..2d84d78d56c --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.integration.test.tsx @@ -0,0 +1,722 @@ +import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; +import * as networking from "@/components/networking"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { toast } from "@/lib/toast"; +import type { Team } from "../key_team_helpers/key_list"; +import { MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES } from "./memberBudgetReset"; +import TeamInfoView, { type TeamData } from "./TeamInfo"; + +const authState = vi.hoisted(() => ({ userRole: "Admin" })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + token: "123", + accessToken: "123", + userId: "user-1", + userEmail: "user@example.com", + userRole: authState.userRole, + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }), +})); + +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + +vi.mock("@/components/networking", () => ({ + serverRootPath: "", + teamInfoCall: vi.fn(), + teamMemberDeleteCall: vi.fn(), + teamMemberAddCall: vi.fn(), + teamMemberUpdateCall: vi.fn(), + teamUpdateCall: vi.fn(), + getGuardrailsList: vi.fn(), + getPoliciesList: vi.fn(), + getPolicyInfoWithGuardrails: vi.fn(), + fetchMCPAccessGroups: vi.fn(), + getTeamPermissionsCall: vi.fn(), + organizationInfoCall: vi.fn(), + getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }), + getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + listMCPTools: vi.fn().mockResolvedValue({ tools: [] }), + vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), +})); + +const { bulkUpdatePOST } = vi.hoisted(() => ({ bulkUpdatePOST: vi.fn() })); +vi.mock("@/lib/http/api", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } }; +}); + +const can = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ + default: (...args: unknown[]) => can(...args), +})); + +vi.mock("@/components/utils/dataUtils", () => ({ + copyToClipboard: vi.fn().mockResolvedValue(true), + formatNumberWithCommas: vi.fn((value: number) => value.toLocaleString()), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ + useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useAllProxyModels: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeam: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + organizationKeys: { all: ["organizations"] }, + useOrganization: vi.fn(), + useOrganizations: vi.fn().mockReturnValue({ data: [], isLoading: false }), +})); + +vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ + useCurrentUser: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn(), +})); + +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + default: ({ + value, + onChange, + }: { + value?: { servers: string[]; accessGroups: string[]; toolsets?: string[] }; + onChange: (next: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; + }) => ( + <> + + + + ), +})); + +vi.mock("@/components/team/TeamMemberTab", () => ({ + default: vi.fn(({ setIsAddMemberModalVisible }) => ( +
+ +
+ )), +})); + +vi.mock("@/components/common_components/user_search_modal", () => ({ + default: vi.fn(({ isVisible, onCancel, onSubmit }) => + isVisible ? ( +
+ + +
+ ) : null, + ), +})); + +vi.mock("@/components/team/EditMembership", () => ({ + default: vi.fn(({ visible, onCancel, onSubmit }) => + visible ? ( +
+ + +
+ ) : null, + ), +})); + +vi.mock("@/components/common_components/DeleteResourceModal", () => ({ + default: vi.fn(({ isOpen, onCancel, onOk }) => + isOpen ? ( +
+ + +
+ ) : null, + ), +})); + +vi.mock("@/components/team/member_permissions", () => ({ + default: vi.fn(() =>
Member Permissions
), +})); + +vi.mock("@/components/common_components/ModelAliasManager", () => ({ + default: vi.fn(({ initialModelAliases, onAliasUpdate }) => ( +
+
{JSON.stringify(initialModelAliases)}
+ + +
+ )), +})); + +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1", access_mcp_server_ids: [] }, + { access_group_id: "ag-2", access_group_name: "Group 2", access_mcp_server_ids: [] }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: ({ value, onChange }: { value?: string[]; onChange?: (next: string[]) => void }) => ( + + ), +})); + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => { + const useKeysResult = { + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + }; + return { useKeys: vi.fn().mockReturnValue(useKeysResult) }; +}); + +vi.mock("../key_team_helpers/filter_helpers", () => ({ + fetchTeamFilterOptions: vi.fn().mockResolvedValue({ + keyAliases: [], + organizationIds: [], + userIds: [], + }), + fetchAllKeyAliases: vi.fn().mockResolvedValue([]), + fetchAllOrganizations: vi.fn().mockResolvedValue([]), +})); + +import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; + +const mockUseAllProxyModels = vi.mocked(useAllProxyModels); +const mockUseKeys = vi.mocked(useKeys); +const mockUseTeam = vi.mocked(useTeam); +const mockUseOrganization = vi.mocked(useOrganization); +const mockUseCurrentUser = vi.mocked(useCurrentUser); +const mockUseMCPServers = vi.mocked(useMCPServers); +const mockUseMCPToolsets = vi.mocked(useMCPToolsets); +const mockUseAccessGroups = vi.mocked(useAccessGroups); +const mockUseUISettings = vi.mocked(useUISettings); + +const createMockTeamData = (overrides = {}) => ({ + team_id: "123", + team_info: { + team_alias: "Test Team", + team_id: "123", + organization_id: null, + admins: ["admin@test.com"], + members: ["user1@test.com"], + members_with_roles: [ + { + user_id: "user1@test.com", + user_email: "user1@test.com", + role: "member", + spend: 0, + budget_id: "budget1", + }, + ], + metadata: {}, + tpm_limit: null, + rpm_limit: null, + max_budget: null, + budget_duration: null, + models: [], + blocked: false, + spend: 0, + max_parallel_requests: null, + budget_reset_at: null, + model_id: null, + litellm_model_table: null, + created_at: "2024-01-01T00:00:00Z", + team_member_budget_table: null, + guardrails: [], + policies: [], + object_permission: null, + ...overrides, + }, + keys: [], + team_memberships: [], +}); + +const seedDefaultMocks = () => { + mockUseAllProxyModels.mockReturnValue({ + data: { data: [] }, + isLoading: false, + } as unknown as ReturnType); + mockUseTeam.mockReturnValue({ + data: undefined, + isLoading: false, + } as unknown as ReturnType); + mockUseOrganization.mockReturnValue({ + data: undefined, + isLoading: false, + } as unknown as ReturnType); + mockUseCurrentUser.mockReturnValue({ + data: { models: [] }, + isLoading: false, + } as unknown as ReturnType); + mockUseMCPServers.mockReturnValue({ data: [], isLoading: false, isError: false } as unknown as ReturnType< + typeof useMCPServers + >); + mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false, isError: false } as unknown as ReturnType< + typeof useMCPToolsets + >); + mockUseAccessGroups.mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1", access_mcp_server_ids: [] }, + { access_group_id: "ag-2", access_group_name: "Group 2", access_mcp_server_ids: [] }, + ], + isLoading: false, + isError: false, + } as unknown as ReturnType); + mockUseUISettings.mockReturnValue({ + data: { values: {} }, + isLoading: false, + } as unknown as ReturnType); + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as unknown as ReturnType); + vi.mocked(useTeamMetadataSchema).mockReturnValue({ + data: [], + isLoading: false, + } as unknown as ReturnType); + + can.mockReturnValue(true); + vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(networking.getTeamPermissionsCall).mockResolvedValue({ + all_available_permissions: [], + team_member_permissions: [], + }); +}; +describe("TeamInfoView - member budget reset prompt", () => { + const props = { + teamId: "123", + onUpdate: vi.fn(), + onClose: vi.fn(), + accessToken: "test-token", + is_team_admin: true, + is_proxy_admin: true, + userModels: ["gpt-4", "gpt-3.5-turbo"], + editTeam: false, + premiumUser: false, + }; + + const customBudgetMembership = ( + userId: string, + maxBudget: number | null = 50, + ): TeamData["team_memberships"][number] => ({ + user_id: userId, + team_id: "123", + budget_id: `budget-${userId}`, + budget_source: "custom", + spend: 0, + total_spend: 0, + litellm_budget_table: { + budget_id: `budget-${userId}`, + soft_budget: null, + max_budget: maxBudget, + max_parallel_requests: null, + tpm_limit: null, + rpm_limit: null, + model_max_budget: null, + budget_duration: null, + budget_reset_at: null, + }, + }); + + const savedTeam: Team = { + team_id: "123", + team_alias: "Test Team", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org-1", + created_at: "2024-01-01T00:00:00Z", + keys: [], + members_with_roles: [], + spend: 0, + }; + + const openEditorWithCustomMembers = async ( + user: ReturnType, + userIds: string[] = ["user-custom"], + maxBudget: number | null = 50, + ) => { + const data = { + ...createMockTeamData({ + team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }), + team_memberships: userIds.map((id) => customBudgetMembership(id, maxBudget)), + } as TeamData; + vi.mocked(networking.teamInfoCall).mockResolvedValue(data); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" }); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByText("Team Member Settings")); + return await screen.findByLabelText("Default Budget (USD)"); + }; + + const submitNewDefault = async (user: ReturnType, input: HTMLElement, value: string) => { + fireEvent.change(input, { target: { value } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + }; + + beforeEach(() => { + seedDefaultMocks(); + bulkUpdatePOST.mockResolvedValue({ data: { data: [] } }); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("holds the save behind the prompt when the default changes while a member has a custom budget", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + + await submitNewDefault(user, input, "20"); + + expect(await screen.findByText("Reset member budgets?")).toBeInTheDocument(); + expect(screen.getByText(/1 member has a custom budget/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Reset to $20.00" })).toBeInTheDocument(); + expect(networking.teamUpdateCall).not.toHaveBeenCalled(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("saves the new default and leaves custom budgets alone on keep", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Keep custom budget" })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + }); + + it("saves the new default then clears each custom budget through the bulk endpoint on reset all", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, ["user-a", "user-b"]); + bulkUpdatePOST.mockResolvedValue({ + data: { + data: [ + { success: true, user_id: "user-a" }, + { success: true, user_id: "user-b" }, + ], + }, + }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(1)); + expect(networking.teamUpdateCall).toHaveBeenCalled(); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(bulkUpdatePOST).toHaveBeenCalledWith("/management/v1/teams/{team_id}/members/bulk_update", { + params: { path: { team_id: "123" } }, + body: { + members: [ + { user_id: "user-a", max_budget_in_team: null }, + { user_id: "user-b", max_budget_in_team: null }, + ], + }, + }); + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Reset 2 member budgets to the team default")); + }); + + it("keeps the prompt mounted while the member-budget reset is in flight", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + const bulk = Promise.withResolvers<{ data: { data: { success: boolean; user_id: string }[] } }>(); + vi.mocked(networking.teamInfoCall).mockImplementationOnce(() => new Promise(() => {})); + bulkUpdatePOST.mockImplementationOnce(() => bulk.promise); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalled()); + expect(screen.getByText("Reset member budgets?")).toBeInTheDocument(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + + bulk.resolve({ data: { data: [{ success: true, user_id: "user-custom" }] } }); + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + }); + + it("offers a retry after a failed reset and never re-saves the team", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + bulkUpdatePOST + .mockRejectedValueOnce(new Error("bulk update down")) + .mockResolvedValueOnce({ data: { data: [{ success: true, user_id: "user-custom" }] } }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + + await waitFor(() => + expect(toast.fromError).toHaveBeenCalledWith("Team updated, but member budgets could not be reset"), + ); + expect(screen.getByText("Reset member budgets?")).toBeInTheDocument(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + await waitFor(() => expect(vi.mocked(networking.teamInfoCall).mock.calls.length).toBeGreaterThan(1)); + + await user.click(screen.getByRole("button", { name: "Retry reset" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(2)); + expect(vi.mocked(networking.teamUpdateCall).mock.calls).toHaveLength(1); + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Reset 1 member budget to the team default")); + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + }); + + it("returns to the prompt without its own toast when the team save fails", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + vi.mocked(networking.teamUpdateCall).mockRejectedValueOnce(new Error("save failed")); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + + expect(await screen.findByRole("button", { name: "Reset to $20.00" })).toBeEnabled(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + }); + + it("reports how many budgets were already reset when a later batch fails", async () => { + const user = userEvent.setup({ delay: null }); + const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `user-${i}`); + const input = await openEditorWithCustomMembers(user, userIds); + bulkUpdatePOST + .mockResolvedValueOnce({ + data: { + data: userIds.slice(0, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES).map((user_id) => ({ success: true, user_id })), + }, + }) + .mockRejectedValueOnce(new Error("second batch down")) + .mockResolvedValueOnce({ + data: { data: [{ success: true, user_id: `user-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}` }] }, + }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" })); + + await waitFor(() => + expect(toast.error).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES} of ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets; the rest could not be reset`, + ), + ); + expect(bulkUpdatePOST).toHaveBeenCalledTimes(2); + + await user.click(screen.getByRole("button", { name: "Retry reset" })); + + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalledTimes(3)); + expect(vi.mocked(bulkUpdatePOST).mock.calls[2][1]).toEqual({ + params: { path: { team_id: "123" } }, + body: { + members: [{ user_id: `user-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`, max_budget_in_team: null }], + }, + }); + expect(vi.mocked(networking.teamUpdateCall).mock.calls).toHaveLength(1); + await waitFor(() => + expect(toast.success).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets to the team default`, + ), + ); + }); + + it("surfaces a failure toast when some members cannot be reset", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, ["user-a", "user-b"]); + bulkUpdatePOST.mockResolvedValue({ + data: { + data: [ + { success: true, user_id: "user-a" }, + { success: false, user_id: "user-b", error: "no such member" }, + ], + }, + }); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset all to $20.00" })); + + await waitFor(() => + expect(toast.error).toHaveBeenCalledWith("Team updated, but 1 member budget could not be reset"), + ); + }); + + it("aborts the save entirely on cancel", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + expect(networking.teamUpdateCall).not.toHaveBeenCalled(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("saves directly when no member carries a custom budget", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, []); + + await submitNewDefault(user, input, "20"); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("saves directly when the default is resubmitted unchanged", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithCustomMembers(user); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(10); + expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("dismisses a pending prompt when the selected team changes", async () => { + const user = userEvent.setup({ delay: null }); + const data = { + ...createMockTeamData({ + team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }), + team_memberships: [customBudgetMembership("user-custom")], + } as TeamData; + vi.mocked(networking.teamInfoCall).mockResolvedValue(data); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" }); + + const { rerender } = renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByText("Team Member Settings")); + const input = await screen.findByLabelText("Default Budget (USD)"); + + await submitNewDefault(user, input, "20"); + expect(await screen.findByText("Reset member budgets?")).toBeInTheDocument(); + + rerender(); + + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); + + it("does not reopen the dialog or refetch the old team when the team changes mid-reset", async () => { + const user = userEvent.setup({ delay: null }); + const data = { + ...createMockTeamData({ + team_member_budget_table: { max_budget: 10, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }), + team_memberships: [customBudgetMembership("user-custom")], + } as TeamData; + vi.mocked(networking.teamInfoCall).mockResolvedValue(data); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" }); + const bulkDone = Promise.withResolvers<{ data: { data: { success: boolean; user_id: string }[] } }>(); + bulkUpdatePOST.mockReturnValue(bulkDone.promise); + + const { rerender } = renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await user.click(screen.getByText("Team Member Settings")); + const input = await screen.findByLabelText("Default Budget (USD)"); + + await submitNewDefault(user, input, "20"); + await user.click(await screen.findByRole("button", { name: "Reset to $20.00" })); + await waitFor(() => expect(bulkUpdatePOST).toHaveBeenCalled()); + + const infoCallsBeforeSwitch = vi.mocked(networking.teamInfoCall).mock.calls.length; + rerender(); + bulkDone.resolve({ data: { data: [{ success: true, user_id: "user-custom" }] } }); + + await waitFor(() => expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument()); + await waitFor(() => + expect( + vi + .mocked(networking.teamInfoCall) + .mock.calls.slice(infoCallsBeforeSwitch) + .some((call) => call[1] === "456"), + ).toBe(true), + ); + const infoCallsAfterSwitch = vi.mocked(networking.teamInfoCall).mock.calls.slice(infoCallsBeforeSwitch); + expect(infoCallsAfterSwitch.every((call) => call[1] === "456")).toBe(true); + expect(screen.queryByRole("button", { name: "Retry reset" })).not.toBeInTheDocument(); + }); + + it("saves directly when a custom member's cap is null and already inherits the default", async () => { + const user = userEvent.setup({ delay: null }); + const input = await openEditorWithCustomMembers(user, ["user-limits-only"], null); + + await submitNewDefault(user, input, "20"); + + await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled()); + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_budget).toBe(20); + expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument(); + expect(bulkUpdatePOST).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 18b1505e342..37f63433d2a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -55,6 +55,12 @@ vi.mock("@/components/networking", () => ({ getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), })); +const { bulkUpdatePOST } = vi.hoisted(() => ({ bulkUpdatePOST: vi.fn() })); +vi.mock("@/lib/http/api", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, fetchClient: { ...actual.fetchClient, POST: bulkUpdatePOST } }; +}); + const can = vi.fn(); vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: (...args: unknown[]) => can(...args), diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index d9e308e9d6f..78ed507216a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -119,6 +119,10 @@ import { } from "./tabVisibilityUtils"; import TeamMembersComponent from "./TeamMemberTab"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; +import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog"; +import { customBudgetMemberUserIds, shouldPromptMemberBudgetReset } from "./memberBudgetReset"; +import { useMemberBudgetReset } from "./useMemberBudgetReset"; +import { fetchClient } from "@/lib/http/api"; const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ "logging", @@ -884,17 +888,41 @@ const TeamInfoView: React.FC = ({ const persistTeamUpdate = async (token: string, updateData: Record) => { await teamUpdateCall(token, updateData); queryClient.invalidateQueries({ queryKey: organizationKeys.all }); - - toast.success("Team settings updated successfully"); setIsEditing(false); - fetchTeamInfo(); }; + const memberBudgetReset = useMemberBudgetReset({ + saveTeam: async (updateData) => { + if (!accessToken) return; + setIsTeamSaving(true); + try { + await persistTeamUpdate(accessToken, updateData); + } finally { + setIsTeamSaving(false); + } + }, + resetMemberBudgets: async (bulkTeamId, userIds) => { + const { data } = await fetchClient.POST("/management/v1/teams/{team_id}/members/bulk_update", { + params: { path: { team_id: bulkTeamId } }, + body: { members: userIds.map((user_id) => ({ user_id, max_budget_in_team: null })) }, + }); + return data?.data ?? []; + }, + refreshTeamData, + }); + + const { dismiss: dismissMemberBudgetReset } = memberBudgetReset; + useEffect(() => { + dismissMemberBudgetReset(); + }, [teamId, dismissMemberBudgetReset]); + const saveTeamAdminSettings = async (changes: TeamAdminSettingsChanges) => { if (!accessToken) return; setIsTeamSaving(true); try { await persistTeamUpdate(accessToken, { team_id: teamId, ...changes }); + toast.success("Team settings updated successfully"); + await fetchTeamInfo(); } catch (error) { console.error("Error updating team:", error); } finally { @@ -1005,8 +1033,10 @@ const TeamInfoView: React.FC = ({ updateData.max_budget = mapEmptyStringToNull(updateData.max_budget); updateData.team_member_budget_duration = values.team_member_budget_duration; - if (values.team_member_budget !== undefined) { - updateData.team_member_budget = Number(values.team_member_budget); + const newTeamMemberBudget = + values.team_member_budget !== undefined ? Number(values.team_member_budget) : undefined; + if (newTeamMemberBudget !== undefined) { + updateData.team_member_budget = newTeamMemberBudget; } if (values.team_member_key_duration !== undefined) { @@ -1152,7 +1182,28 @@ const TeamInfoView: React.FC = ({ } } + const customBudgetUserIds = customBudgetMemberUserIds(teamData?.team_memberships ?? []); + if ( + newTeamMemberBudget !== undefined && + shouldPromptMemberBudgetReset( + newTeamMemberBudget, + info.team_member_budget_table?.max_budget, + customBudgetUserIds, + ) + ) { + const pendingReset = { + teamId, + updateData, + userIds: customBudgetUserIds, + newBudget: newTeamMemberBudget, + }; + memberBudgetReset.prompt(pendingReset); + return; + } + await persistTeamUpdate(accessToken, updateData); + toast.success("Team settings updated successfully"); + await fetchTeamInfo(); } catch (error) { console.error("Error updating team:", error); } finally { @@ -2426,6 +2477,14 @@ const TeamInfoView: React.FC = ({ onOk={handleDeleteConfirm} confirmLoading={isDeleting} /> + + ); }; diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 8652ffa7de2..bd7398bc6c8 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -682,6 +682,44 @@ describe("TeamMembersComponent", () => { expect(inheritedRow).toHaveTextContent("$25.00"); }); + it("caps a Custom member at the team default when the private row has no budget limit", () => { + const base = createMockTeamData(); + renderTab( + createMockTeamData({ + team_info: { + ...base.team_info, + team_member_budget_table: { max_budget: 20, budget_duration: null, tpm_limit: null, rpm_limit: null }, + }, + team_memberships: [ + { + user_id: "user2@test.com", + team_id: "team-123", + budget_id: "budget2", + budget_source: "custom", + spend: 0, + total_spend: 0, + litellm_budget_table: { + budget_id: "budget3", + soft_budget: null, + max_budget: null, + max_parallel_requests: null, + tpm_limit: null, + rpm_limit: 100, + model_max_budget: null, + budget_duration: null, + budget_reset_at: null, + }, + }, + ], + }), + ); + + const row = screen.getByRole("row", { name: /user2@test\.com/ }); + expect(within(row).getByTestId("member-budget-source")).toHaveTextContent("Custom"); + expect(row).toHaveTextContent("$20.00"); + expect(row).not.toHaveTextContent("Unlimited"); + }); + it("shows no source label for a member with neither a custom nor a team budget", () => { renderTab(createMockTeamData({ team_memberships: [] })); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 660416504fe..e8576ce6dfe 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -106,10 +106,7 @@ export default function TeamMemberTab({ const getUserBudget = (userId: string | null): number | null => { if (!userId) return null; const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return ( - membership?.litellm_budget_table?.max_budget ?? - (membership?.budget_source === "team_default" ? teamDefaultBudget : null) - ); + return membership?.litellm_budget_table?.max_budget ?? teamDefaultBudget; }; // Helper function to get rate limits for a user diff --git a/ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts b/ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts new file mode 100644 index 00000000000..e673840203f --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/memberBudgetReset.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, it } from "vitest"; +import { + chunk, + customBudgetMemberUserIds, + MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES, + pluralize, + shouldPromptMemberBudgetReset, +} from "./memberBudgetReset"; + +describe("customBudgetMemberUserIds", () => { + const customRow = (user_id: string, max_budget: number | null = 50) => ({ + user_id, + budget_source: "custom" as const, + litellm_budget_table: { max_budget }, + }); + + it("returns only members whose budget_source is custom with a private cap", () => { + const memberships = [ + customRow("u-custom"), + { user_id: "u-default", budget_source: "team_default" as const }, + { user_id: "u-none", budget_source: "none" as const }, + ]; + + expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]); + }); + + it("skips a custom row whose cap is null since it already inherits the default", () => { + const memberships = [customRow("u-rate-limits-only", null), customRow("u-custom")]; + + expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]); + }); + + it("skips a custom row that has no budget table at all", () => { + const memberships = [ + { user_id: "u-no-table", budget_source: "custom" as const, litellm_budget_table: null }, + customRow("u-custom"), + ]; + + expect(customBudgetMemberUserIds(memberships)).toEqual(["u-custom"]); + }); +}); + +describe("shouldPromptMemberBudgetReset", () => { + it("prompts when the default changes and custom-budget members exist", () => { + expect(shouldPromptMemberBudgetReset(10, 5, ["u-1"])).toBe(true); + }); + + it("prompts when a default is set for the first time", () => { + expect(shouldPromptMemberBudgetReset(10, null, ["u-1"])).toBe(true); + expect(shouldPromptMemberBudgetReset(10, undefined, ["u-1"])).toBe(true); + }); + + it("does not prompt when the submitted budget is unchanged", () => { + expect(shouldPromptMemberBudgetReset(10, 10, ["u-1"])).toBe(false); + }); + + it("does not prompt when the budget is cleared or zero", () => { + expect(shouldPromptMemberBudgetReset(undefined, 10, ["u-1"])).toBe(false); + expect(shouldPromptMemberBudgetReset(0, 10, ["u-1"])).toBe(false); + }); + + it("does not prompt when no member has a custom budget", () => { + expect(shouldPromptMemberBudgetReset(10, 5, [])).toBe(false); + }); +}); + +describe("chunk", () => { + it("splits selections larger than the bulk endpoint limit", () => { + const ids = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`); + + const chunks = chunk(ids, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES); + + expect(chunks).toHaveLength(2); + expect(chunks[0]).toHaveLength(MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES); + expect(chunks[1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]); + }); + + it("keeps a selection under the limit in a single chunk", () => { + expect(chunk(["u-1", "u-2"], MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES)).toEqual([["u-1", "u-2"]]); + }); + + it("returns no chunks for an empty selection", () => { + expect(chunk([], MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES)).toEqual([]); + }); + + it("returns no chunks for a non-positive size", () => { + expect(chunk(["u-1"], 0)).toEqual([]); + expect(chunk(["u-1"], -3)).toEqual([]); + }); +}); + +describe("pluralize", () => { + it("uses the singular form for exactly one", () => { + expect(pluralize(1, "budget", "budgets")).toBe("budget"); + }); + + it("uses the plural form for zero and for more than one", () => { + expect(pluralize(0, "budget", "budgets")).toBe("budgets"); + expect(pluralize(3, "budget", "budgets")).toBe("budgets"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/memberBudgetReset.ts b/ui/litellm-dashboard/src/components/team/memberBudgetReset.ts new file mode 100644 index 00000000000..2b24a5ebbb0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/memberBudgetReset.ts @@ -0,0 +1,38 @@ +import type { components } from "@/lib/http/schema"; +import type { TeamMembership } from "./TeamInfo"; + +export const MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES = 500; + +export type TeamUpdatePayload = components["schemas"]["UpdateTeamRequest"]; + +export interface MemberBudgetResetPending { + readonly teamId: string; + readonly updateData: TeamUpdatePayload; + readonly userIds: readonly string[]; + readonly newBudget: number; +} + +type MembershipBudgetRow = Pick & { + litellm_budget_table?: Pick | null; +}; + +export const customBudgetMemberUserIds = (memberships: readonly MembershipBudgetRow[]): string[] => + memberships + .filter((m) => m.budget_source === "custom" && m.litellm_budget_table?.max_budget != null) + .map((m) => m.user_id); + +export const shouldPromptMemberBudgetReset = ( + nextBudget: number | undefined, + previousBudget: number | null | undefined, + customBudgetUserIds: string[], +): boolean => { + const budgetChanged = typeof nextBudget === "number" && nextBudget > 0 && nextBudget !== previousBudget; + return budgetChanged && customBudgetUserIds.length > 0; +}; + +export const chunk = (items: readonly T[], size: number): T[][] => { + if (size <= 0) return []; + return Array.from({ length: Math.ceil(items.length / size) }, (_, i) => items.slice(i * size, i * size + size)); +}; + +export const pluralize = (count: number, singular: string, plural: string): string => (count === 1 ? singular : plural); diff --git a/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx new file mode 100644 index 00000000000..71318b635a3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.test.tsx @@ -0,0 +1,345 @@ +import { act, renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { toast } from "@/lib/toast"; +import { MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES } from "./memberBudgetReset"; +import { useMemberBudgetReset, type MemberBudgetResetGateway } from "./useMemberBudgetReset"; + +const buildGateway = () => ({ + saveTeam: vi.fn(async (_updateData: Record) => {}), + resetMemberBudgets: vi.fn(async (_teamId: string, userIds: readonly string[]) => + userIds.map((user_id) => ({ success: true, user_id })), + ), + refreshTeamData: vi.fn(async () => {}), +}); + +const pendingFor = (userIds: string[]) => ({ + teamId: "team-123", + updateData: { team_id: "team-123", team_member_budget: 20 }, + userIds, + newBudget: 20, +}); + +const renderReset = (gateway: MemberBudgetResetGateway) => renderHook(() => useMemberBudgetReset(gateway)); + +describe("useMemberBudgetReset", () => { + afterEach(() => { + vi.clearAllMocks(); + }); + + it("starts idle and prompts with the pending update", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + expect(result.current.state.phase).toBe("idle"); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + + expect(result.current.state).toEqual({ + phase: "prompting", + pending: pendingFor(["u-1"]), + }); + expect(gateway.saveTeam).not.toHaveBeenCalled(); + }); + + it("saves the team once, resets every member, and refreshes on success", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1", "u-2"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(gateway.saveTeam).toHaveBeenCalledTimes(1); + expect(gateway.saveTeam).toHaveBeenCalledWith({ team_id: "team-123", team_member_budget: 20 }); + expect(gateway.resetMemberBudgets).toHaveBeenCalledWith("team-123", ["u-1", "u-2"]); + expect(toast.success).toHaveBeenCalledWith("Reset 2 member budgets to the team default"); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + expect(result.current.state.phase).toBe("idle"); + }); + + it("resets members in batches no larger than the bulk endpoint limit", async () => { + const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`); + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(userIds)); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(gateway.resetMemberBudgets).toHaveBeenCalledTimes(2); + expect(gateway.resetMemberBudgets.mock.calls[0][0]).toBe("team-123"); + expect(gateway.resetMemberBudgets.mock.calls[0][1]).toHaveLength(MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES); + expect(gateway.resetMemberBudgets.mock.calls[1][1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]); + expect(result.current.state.phase).toBe("idle"); + }); + + it("resets against the team carried by the pending update", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + const pending = { + teamId: "team-999", + updateData: { team_id: "team-999", team_member_budget: 20 }, + userIds: ["u-1"], + newBudget: 20, + }; + await act(async () => { + result.current.prompt(pending); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(gateway.resetMemberBudgets).toHaveBeenCalledWith("team-999", ["u-1"]); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalledWith("team-123", expect.anything()); + }); + + it("returns to prompting without its own toast when the team save fails", async () => { + const gateway = buildGateway(); + gateway.saveTeam.mockRejectedValueOnce(new Error("team update failed")); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(result.current.state.phase).toBe("prompting"); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(toast.success).not.toHaveBeenCalled(); + }); + + it("retries only the unsent members when a later batch fails, without re-saving the team", async () => { + const userIds = Array.from({ length: MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1 }, (_, i) => `u-${i}`); + const gateway = buildGateway(); + gateway.resetMemberBudgets + .mockResolvedValueOnce( + userIds.slice(0, MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES).map((user_id) => ({ success: true, user_id })), + ) + .mockRejectedValueOnce(new Error("second batch failed")); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(userIds)); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(result.current.state.phase).toBe("resetFailed"); + expect(toast.error).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES} of ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets; the rest could not be reset`, + ); + + gateway.resetMemberBudgets.mockResolvedValueOnce([ + { success: true, user_id: `u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}` }, + ]); + await act(async () => { + await result.current.retry(); + }); + + expect(gateway.saveTeam).toHaveBeenCalledTimes(1); + expect(gateway.resetMemberBudgets).toHaveBeenCalledTimes(3); + expect(gateway.resetMemberBudgets.mock.calls[2][0]).toBe("team-123"); + expect(gateway.resetMemberBudgets.mock.calls[2][1]).toEqual([`u-${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES}`]); + expect(toast.success).toHaveBeenCalledWith( + `Reset ${MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES + 1} member budgets to the team default`, + ); + expect(result.current.state.phase).toBe("idle"); + }); + + it("reports partial progress when the first batch fails", async () => { + const gateway = buildGateway(); + gateway.resetMemberBudgets.mockRejectedValueOnce(new Error("batch failed")); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1", "u-2"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(result.current.state.phase).toBe("resetFailed"); + expect(toast.fromError).toHaveBeenCalledWith("Team updated, but member budgets could not be reset"); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + }); + + it("reports members the backend could not reset and closes", async () => { + const gateway = buildGateway(); + gateway.resetMemberBudgets.mockResolvedValueOnce([ + { success: true, user_id: "u-1" }, + { success: false, user_id: "u-2" }, + ]); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1", "u-2"])); + }); + await act(async () => { + await result.current.reset(); + }); + + expect(toast.error).toHaveBeenCalledWith("Team updated, but 1 member budget could not be reset"); + expect(toast.success).not.toHaveBeenCalled(); + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + }); + + it("saves the team once and closes on keep custom", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + await act(async () => { + await result.current.keepCustom(); + }); + + expect(gateway.saveTeam).toHaveBeenCalledTimes(1); + expect(toast.success).toHaveBeenCalledWith("Team settings updated successfully"); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + expect(gateway.refreshTeamData).toHaveBeenCalledTimes(1); + expect(result.current.state.phase).toBe("idle"); + }); + + it("abandons the pending update on dismiss", async () => { + const gateway = buildGateway(); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + await act(async () => { + result.current.dismiss(); + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.saveTeam).not.toHaveBeenCalled(); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + }); + + it("skips the refresh and stays idle when dismissed while a reset is in flight", async () => { + const gateway = buildGateway(); + const bulkDone = Promise.withResolvers<{ success: boolean; user_id: string }[]>(); + gateway.resetMemberBudgets.mockReturnValueOnce(bulkDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let resetPromise = Promise.resolve(); + await act(async () => { + resetPromise = result.current.reset(); + }); + expect(result.current.state.phase).toBe("resetting"); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + bulkDone.resolve([{ success: true, user_id: "u-1" }]); + await resetPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.success).not.toHaveBeenCalled(); + }); + + it("does not return to prompting when the team save fails after a dismiss", async () => { + const gateway = buildGateway(); + const saveDone = Promise.withResolvers(); + gateway.saveTeam.mockReturnValueOnce(saveDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let resetPromise = Promise.resolve(); + await act(async () => { + resetPromise = result.current.reset(); + }); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + saveDone.reject(new Error("team update failed")); + await resetPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.resetMemberBudgets).not.toHaveBeenCalled(); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + }); + + it("stays idle without a refresh when the reset fails after a dismiss", async () => { + const gateway = buildGateway(); + const bulkDone = Promise.withResolvers<{ success: boolean; user_id: string }[]>(); + gateway.resetMemberBudgets.mockReturnValueOnce(bulkDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let resetPromise = Promise.resolve(); + await act(async () => { + resetPromise = result.current.reset(); + }); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + bulkDone.reject(new Error("bulk failed")); + await resetPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.error).not.toHaveBeenCalled(); + expect(toast.fromError).not.toHaveBeenCalled(); + }); + + it("skips the refresh when dismissed while keep-custom is saving", async () => { + const gateway = buildGateway(); + const saveDone = Promise.withResolvers(); + gateway.saveTeam.mockReturnValueOnce(saveDone.promise); + const { result } = renderReset(gateway); + + await act(async () => { + result.current.prompt(pendingFor(["u-1"])); + }); + let keepPromise = Promise.resolve(); + await act(async () => { + keepPromise = result.current.keepCustom(); + }); + + act(() => { + result.current.dismiss(); + }); + await act(async () => { + saveDone.resolve(); + await keepPromise; + }); + + expect(result.current.state.phase).toBe("idle"); + expect(gateway.refreshTeamData).not.toHaveBeenCalled(); + expect(toast.success).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts new file mode 100644 index 00000000000..f0ff832b34a --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/useMemberBudgetReset.ts @@ -0,0 +1,134 @@ +import { useCallback, useRef, useState } from "react"; +import type { components } from "@/lib/http/schema"; +import { toast } from "@/lib/toast"; +import { + chunk, + MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES, + pluralize, + type MemberBudgetResetPending, + type TeamUpdatePayload, +} from "./memberBudgetReset"; + +export type MemberBudgetBulkResult = components["schemas"]["TeamMemberBudgetUpdateResult"]; + +export type MemberBudgetResetState = + | { phase: "idle" } + | { phase: "prompting"; pending: MemberBudgetResetPending } + | { phase: "resetting"; pending: MemberBudgetResetPending; attempted: number } + | { phase: "resetFailed"; pending: MemberBudgetResetPending; attempted: number }; + +export interface MemberBudgetResetGateway { + saveTeam: (updateData: TeamUpdatePayload) => Promise; + resetMemberBudgets: (teamId: string, userIds: readonly string[]) => Promise; + refreshTeamData: () => Promise; +} + +type ResetRun = + | { ok: true; attempted: number; results: readonly MemberBudgetBulkResult[] } + | { ok: false; attempted: number; error: unknown }; + +export const useMemberBudgetReset = (gateway: MemberBudgetResetGateway) => { + const [state, setState] = useState({ phase: "idle" }); + const activeRun = useRef(null); + const isCurrent = (run: object) => activeRun.current === run; + + const runReset = async (pending: MemberBudgetResetPending, attempted: number, run: object) => { + const { resetMemberBudgets, refreshTeamData } = gateway; + + const runChunks = ( + chunks: readonly (readonly string[])[], + attemptedSoFar: number, + results: readonly MemberBudgetBulkResult[], + ): Promise => { + const [ids, ...rest] = chunks; + if (ids === undefined) return Promise.resolve({ ok: true, attempted: attemptedSoFar, results }); + return resetMemberBudgets(pending.teamId, ids).then( + (batch) => runChunks(rest, attemptedSoFar + ids.length, [...results, ...batch]), + (error: unknown) => ({ ok: false as const, attempted: attemptedSoFar, error }), + ); + }; + + const outcome = await runChunks( + chunk(pending.userIds.slice(attempted), MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES), + attempted, + [], + ); + + if (!isCurrent(run)) return; + if (!outcome.ok) { + console.error("Error resetting member budgets:", outcome.error); + const total = pending.userIds.length; + if (outcome.attempted > 0) { + toast.error( + `Reset ${outcome.attempted} of ${total} member ${pluralize(total, "budget", "budgets")}; the rest could not be reset`, + ); + } else { + toast.fromError("Team updated, but member budgets could not be reset"); + } + setState({ phase: "resetFailed", pending, attempted: outcome.attempted }); + await refreshTeamData(); + return; + } + const failed = outcome.results.filter((r) => !r.success); + if (failed.length > 0) { + toast.error( + `Team updated, but ${failed.length} member ${pluralize(failed.length, "budget", "budgets")} could not be reset`, + ); + } else { + toast.success( + `Reset ${outcome.attempted} member ${pluralize(outcome.attempted, "budget", "budgets")} to the team default`, + ); + } + setState({ phase: "idle" }); + await refreshTeamData(); + }; + + const prompt = (pending: MemberBudgetResetPending) => setState({ phase: "prompting", pending }); + + const reset = async () => { + if (state.phase !== "prompting") return; + const { pending } = state; + const run = {}; + activeRun.current = run; + setState({ phase: "resetting", pending, attempted: 0 }); + try { + await gateway.saveTeam(pending.updateData); + } catch (error) { + console.error("Error updating team:", error); + if (isCurrent(run)) setState({ phase: "prompting", pending }); + return; + } + await runReset(pending, 0, run); + }; + + const retry = async () => { + if (state.phase !== "resetFailed") return; + const { pending, attempted } = state; + const run = {}; + activeRun.current = run; + setState({ phase: "resetting", pending, attempted }); + await runReset(pending, attempted, run); + }; + + const keepCustom = async () => { + if (state.phase !== "prompting") return; + const { pending } = state; + const run = {}; + activeRun.current = run; + setState({ phase: "idle" }); + try { + await gateway.saveTeam(pending.updateData); + if (isCurrent(run)) toast.success("Team settings updated successfully"); + } catch (error) { + console.error("Error updating team:", error); + } + if (isCurrent(run)) await gateway.refreshTeamData(); + }; + + const dismiss = useCallback(() => { + activeRun.current = null; + setState({ phase: "idle" }); + }, []); + + return { state, prompt, reset, retry, keepCustom, dismiss }; +}; From e450f2d54c113a962c38e32cf3d0d29341991c3e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:31:43 -0700 Subject: [PATCH 4/6] fix(router): serve Responses turns from a sibling when the encrypted content origin has no boundary peer (#43015) * test(integration): reproduce encrypted_content_affinity 503 when origin has no boundary peer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): keep encrypted content affinity turn one out of the response cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): degrade encrypted_content_affinity when the origin has no encryption-boundary peer A routed-group candidate that is currently unavailable and shares its (api_base, api_key) with no healthy deployment used to raise a proxy-level 503/429 from _unavailable_origin_error, even though healthy siblings in the same model group could still serve the turn. Strip the encrypted reasoning and dispatch to the healthy pool instead, matching the existing cross-group behavior, and log a warning naming the origin model_id and routed group Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): bound the degraded-affinity log marker and assert the strip on the wire Address review findings: restore num_retries alongside optional_pre_call_checks in the integration test teardown, record scenario request bodies on the scripted upstream so the tests can assert no encrypted reasoning reaches the sibling, and truncate the client-supplied model_id in the degraded-dispatch warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tests): only record JSON bodies on the scripted upstream Multipart uploads to scripted POST routes have no JSON body, so gate the observation recording on the request content-type Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): make encrypted_content_affinity runtime-toggleable so /config/update can turn it off --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router.py | 8 +- .../encrypted_content_affinity_check.py | 166 ++---------- tests/integration/_support/upstream.py | 8 + ...test_encrypted_content_affinity_no_peer.py | 240 ++++++++++++++++++ .../test_router_retry_policy_update.py | 58 +++++ .../test_encrypted_content_affinity_check.py | 237 ++++++++++------- 6 files changed, 471 insertions(+), 246 deletions(-) create mode 100644 tests/integration/routing/test_encrypted_content_affinity_no_peer.py diff --git a/litellm/router.py b/litellm/router.py index da92eef9102..56d27fedf28 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -227,6 +227,9 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, ) +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, +) from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( build_io_token_rate_limit_headers, deployment_has_io_token_limits, @@ -438,6 +441,7 @@ _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = M { "prompt_caching": PromptCachingDeploymentCheck, "enforce_model_rate_limits": ModelRateLimitingCheck, + "encrypted_content_affinity": EncryptedContentAffinityCheck, } ) @@ -2207,10 +2211,6 @@ class Router: ) def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, - ) - def _move_before_deployment_affinity( callback_list: list[Any], callback_to_move: EncryptedContentAffinityCheck, diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index e5e40d7d6f5..cf58f3b3d3c 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -36,17 +36,10 @@ Safe to enable globally: - No cache required. """ -import time from collections.abc import Iterator, Mapping -from typing import TYPE_CHECKING, Final, Optional, Protocol, cast - -import httpx +from typing import TYPE_CHECKING, Final, Optional, cast from litellm._logging import verbose_router_logger -from litellm.exceptions import ( - RateLimitError, - ServiceUnavailableError, -) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -55,7 +48,6 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_encrypted_reasoning_from_messages, ) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues from litellm.types.router import Deployment @@ -63,14 +55,6 @@ if TYPE_CHECKING: from litellm.router import Router -class _SupportsActiveCooldowns(Protocol): - """Cooldown-cache handle: this check only reads back the currently active cooldowns.""" - - async def async_get_active_cooldowns( - self, model_ids: list[str], parent_otel_span: Span | None - ) -> list[tuple[str, CooldownCacheValue]]: ... - - class EncryptedContentAffinityCheck(CustomLogger): """ Routes follow-up Responses API requests to the deployment that produced @@ -194,23 +178,6 @@ class EncryptedContentAffinityCheck(CustomLogger): return deployment return None - @staticmethod - def _request_team_id(request_kwargs: Mapping[str, object]) -> str | None: - containers: Final = (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - team_ids: Final = (c.get("user_api_key_team_id") for c in containers if isinstance(c, Mapping)) - return next((tid for tid in team_ids if isinstance(tid, str)), None) - - def _routed_group_candidate_model_ids(self, request_kwargs: Mapping[str, object], model: str) -> frozenset[str]: - """ - Deployment ids that could serve this turn's routed ``model``, as the router - resolves a route (model_group_alias / routing group / model_name / team / - pattern). Delegates to the router so the full precedence is not re-derived here - and no deployment ids are written into request kwargs bound for the provider. - """ - if self.router is None: - return frozenset() - return self.router.get_candidate_model_ids_for_route(model=model, team_id=self._request_team_id(request_kwargs)) - @staticmethod def _encryption_boundary_key( litellm_params: object, @@ -262,9 +229,6 @@ class EncryptedContentAffinityCheck(CustomLogger): Deployments in ``healthy_deployments`` sharing the originating deployment's ``(api_base, api_key)``, alongside the originating deployment object (or ``None`` if it was removed / router unavailable). - Returns ``([], originating_or_None)`` when no boundary match exists, - so the caller can reuse the looked-up ``originating`` rather than - re-querying the router. """ if self.router is None: return [], None @@ -294,18 +258,12 @@ class EncryptedContentAffinityCheck(CustomLogger): """ If the request ``input`` contains litellm-encoded item IDs, or its Anthropic ``messages`` replay a bridge-tagged thinking block, decode the embedded - ``model_id`` and pin the request to that deployment. Raises - ``RateLimitError`` / ``ServiceUnavailableError`` when the originating - deployment is a member of the routed model group but currently unavailable - and no encryption-boundary peer exists, rather than dispatching a doomed - request to a non-peer deployment. When the origin is not a member of the - routed group (an auto-router tier change, a model switch with no peer, a - removed deployment, or an unknown/forged marker), the encrypted reasoning is - stripped and the request dispatches with its readable history instead. The - 429/503 split mirrors the originating cooldown's status: - a 429-induced cooldown surfaces as 429 (with ``Retry-After`` set to the - remaining cooldown window) so OpenAI-compatible clients back off and - retry after the deployment is eligible again. + ``model_id`` and pin the request to that deployment. When the origin cannot + serve this turn and no encryption-boundary peer is configured (it is + unhealthy, the request was routed to a different group by an auto-router tier + change or model switch, or the marker is removed/unknown/forged), the + encrypted reasoning is stripped and the request dispatches to the healthy + pool with its readable history instead of failing. """ request_kwargs = request_kwargs or {} typed_healthy_deployments: Final = cast(list[dict], healthy_deployments) @@ -348,7 +306,7 @@ class EncryptedContentAffinityCheck(CustomLogger): return [deployment] # Follow-up switched model_name (LIT-2531): pin by Azure resource instead. - boundary_matches, originating = self._find_deployments_on_same_encryption_boundary( + boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary( healthy_deployments=typed_healthy_deployments, model_id=model_id, ) @@ -362,101 +320,17 @@ class EncryptedContentAffinityCheck(CustomLogger): request_kwargs["_encrypted_content_affinity_pinned"] = True return boundary_matches - # The origin cannot serve this turn's routed group and no peer shares the boundary, so its + # The origin cannot serve this turn and no peer shares its encryption boundary, so its # encrypted reasoning can never decrypt here. Strip it, keep the readable history, and dispatch - # to the routed group instead of failing. Membership is tested by deployment id against the set - # the router actually resolved for this route, not by model-group name, so an alias, a - # provider-qualified spelling, a team-public name, or a pattern route of the same group is not - # mistaken for a tier change. An unknown origin (a removed deployment, or a forged marker) is - # treated the same as a cross-group one, which also denies an authenticated caller a - # deployment-id existence oracle: a real cross-group id and a nonexistent id both strip and - # dispatch rather than returning distinguishable responses. Only a genuine same-group member - # that is currently unavailable falls through to the fail-fast, preserving the cooldown contract. - routed_group_model_ids: Final = ( - self._routed_group_candidate_model_ids(request_kwargs, model) if originating is not None else frozenset() + # to the healthy pool instead of failing the request. This also denies an authenticated caller a + # deployment-id existence oracle: a same-group id, a cross-group id, a removed id and a forged + # marker all strip and dispatch rather than returning distinguishable responses. + verbose_router_logger.warning( + "EncryptedContentAffinityCheck: model_id=%s cannot serve group %s and no deployment on the same " + "encryption boundary is configured; forwarding without its encrypted reasoning", + model_id[:64], + model, ) - if str(model_id) not in routed_group_model_ids: - verbose_router_logger.debug( - "EncryptedContentAffinityCheck: model_id=%s is not a candidate for the routed group %s; " - "forwarding without its encrypted reasoning", - model_id, - model, - ) - ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) - strip_encrypted_reasoning_from_messages(anthropic_messages) - return typed_healthy_deployments - - # The origin is a member of the routed group but currently unavailable (cooled down); fail fast - # rather than dispatching to a non-peer, which would guarantee an upstream 400. - raise await self._unavailable_origin_error( - model=model, - model_id=model_id, - parent_otel_span=parent_otel_span, - ) - - async def _unavailable_origin_error( - self, - model: str, - model_id: str, - parent_otel_span: Span | None, - ) -> Exception: - # Public error messages intentionally omit the originating ``model_id`` so - # an authenticated caller forging encrypted-content markers cannot use the - # error surface to enumerate which deployment IDs exist on this router. - cooldown: Final = await self._get_origin_cooldown(model_id=model_id, parent_otel_span=parent_otel_span) - - if cooldown is not None and str(cooldown.get("status_code")) == "429": - retry_after: Final = self._cooldown_seconds_remaining(cooldown) - return RateLimitError( - message=( - "The deployment that produced this encrypted_content is " - f"rate-limited (cooling down for ~{retry_after}s), and no " - "deployment on the same encryption boundary is configured. " - "Retry after the Retry-After window or configure a deployment " - "with the same (api_base, api_key)." - ), - llm_provider="", - model=model, - response=httpx.Response( - status_code=429, - headers={"retry-after": str(retry_after)}, - request=httpx.Request("POST", "https://litellm.ai/"), - ), - ) - - return ServiceUnavailableError( - message=( - "The deployment that produced this encrypted_content is " - "currently unavailable (likely cooled down), and no deployment " - "on the same encryption boundary is configured. Retry later or " - "configure a deployment with the same (api_base, api_key)." - ), - llm_provider="", - model=model, - ) - - async def _get_origin_cooldown( - self, - model_id: str, - parent_otel_span: Span | None, - ) -> CooldownCacheValue | None: - if self.router is None: - return None - cooldown_cache: Final[_SupportsActiveCooldowns | None] = getattr(self.router, "cooldown_cache", None) - if cooldown_cache is None: - return None - try: - active: Final = await cooldown_cache.async_get_active_cooldowns( - model_ids=[model_id], parent_otel_span=parent_otel_span - ) - except Exception: - return None - for cached_model_id, value in active: - if cached_model_id == model_id: - return value - return None - - @staticmethod - def _cooldown_seconds_remaining(cooldown: CooldownCacheValue) -> int: - remaining = float(cooldown.get("timestamp", 0.0)) + float(cooldown.get("cooldown_time", 0.0)) - time.time() - return max(1, int(remaining)) + ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) + strip_encrypted_reasoning_from_messages(anthropic_messages) + return typed_healthy_deployments diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index eea539643d4..e645126032b 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -239,6 +239,14 @@ class Provider: response: Final = self.scenario_store.get(scenario_id) if response is None: return JSONResponse({"error": "Unknown scenario"}, status_code=404) + if request.method == "POST" and "json" in request.headers.get("content-type", ""): + raw_body: Final = await request.body() + if raw_body: + body: Final = JSON_OBJECT.validate_json(raw_body) + if isinstance(body, dict): + self.observations.put( + Observation(request.url.path, request.headers.get("authorization", ""), body) + ) if isinstance(response, RoutedResponse): route_key: Final = f"{request.method} /{'/'.join(segments[1:])}" route: Final = next( diff --git a/tests/integration/routing/test_encrypted_content_affinity_no_peer.py b/tests/integration/routing/test_encrypted_content_affinity_no_peer.py new file mode 100644 index 00000000000..1025278aa10 --- /dev/null +++ b/tests/integration/routing/test_encrypted_content_affinity_no_peer.py @@ -0,0 +1,240 @@ +"""Encrypted-content affinity when the origin deployment has no encryption-boundary peer. + +A multi-region model group has several deployments sharing one upstream api_key but a +distinct api_base each, with ``optional_pre_call_checks: [encrypted_content_affinity]`` +and ``disable_cooldowns: true`` (the integration proxy config). A follow-up +``POST /v1/responses`` that replays a reasoning item must never fail at the proxy while +sibling deployments in the same group are healthy: when the origin cannot serve the turn +its encrypted reasoning should be stripped and the request dispatched to a sibling. +""" + +import json +import uuid +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse +from pydantic import JsonValue + +AFFINITY_CHECK: Final = "encrypted_content_affinity" +PROVIDER_MODEL: Final = "openai/gpt-5" +PROVIDER_KEY: Final = "integration-provider-key" +DEPLOYMENT_COUNT: Final = 3 + + +def _responses_payload() -> dict[str, JsonValue]: + return { + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5-scripted", + "output": [ + { + "type": "reasoning", + "id": "rs_$REQUEST_ID", + "summary": [], + "encrypted_content": "ZHNra2RrZA==", + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted answer"}], + }, + ], + "usage": {"input_tokens": 5, "output_tokens": 7, "total_tokens": 12}, + } + + +def _enable_affinity_check(scenario: Scenario) -> None: + gateway: Final = scenario.gateway + settings: Final = object_value(gateway.get("/router/settings")["current_values"]) + current: Final = settings.get("optional_pre_call_checks") + original: Final = list(current) if isinstance(current, list) else [] + original_retries: Final = settings.get("num_retries") + scenario.cleanups.callback( + lambda: gateway.post( + "/config/update", + {"router_settings": {"optional_pre_call_checks": original, "num_retries": original_retries}}, + ) + ) + gateway.post( + "/config/update", + {"router_settings": {"optional_pre_call_checks": [*original, AFFINITY_CHECK], "num_retries": 0}}, + ) + + +def _last_group_request_body(gateway: Gateway, group: str) -> dict[str, JsonValue]: + """The request body the upstream last saw for this group, to prove the strip reached the wire.""" + with httpx.Client(base_url=gateway.upstream_url, trust_env=False) as upstream: + requests: Final = object_value(upstream.get("/__observations").json()).get("requests") + assert isinstance(requests, list) + bodies: Final = [ + object_value(request).get("body") + for request in requests + if str(object_value(request).get("path")).startswith(f"/{group}-") + ] + assert bodies, f"upstream saw no requests for {group}: {requests}" + body: Final = bodies[-1] + assert isinstance(body, dict), f"upstream request body is not an object: {body}" + return body + + +def _assert_no_encrypted_reasoning_reached_upstream(gateway: Gateway, group: str) -> None: + body: Final = _last_group_request_body(gateway, group) + items: Final = body.get("input") + assert isinstance(items, list), f"upstream request carried no input list: {body}" + assert not any( + isinstance(item, dict) and ("encrypted_content" in item or "litellm_enc" in json.dumps(item)) for item in items + ), f"undecryptable reasoning reached the sibling: {body}" + + +def _multi_region_group(scenario: Scenario) -> tuple[str, tuple[str, ...]]: + """Three deployments in one model group: one shared api_key, a distinct api_base each.""" + gateway: Final = scenario.gateway + group: Final = f"enc-affinity-{uuid.uuid4().hex}" + handles: Final = tuple( + register_scenario( + f"{group}-{index}", + JsonResponse(content_type="application/json", body=_responses_payload()), + ) + for index in range(DEPLOYMENT_COUNT) + ) + for handle in handles: + scenario.cleanups.callback(delete_scenario, handle) + + def delete_model_if_present(model_id: str) -> None: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + if any(object_value(object_value(entry)["model_info"])["id"] == model_id for entry in entries): + scenario.delete_model(model_id) + + deployment_ids: Final = tuple( + string_value( + object_value( + gateway.post( + "/model/new", + { + "model_name": group, + "litellm_params": { + "model": PROVIDER_MODEL, + "api_key": PROVIDER_KEY, + "api_base": handle.api_base(), + }, + "model_info": {}, + }, + )["model_info"] + )["id"] + ) + for handle in handles + ) + for model_id in deployment_ids: + scenario.cleanups.callback(delete_model_if_present, model_id) + return group, deployment_ids + + +def _user_message(text: str) -> dict[str, JsonValue]: + return {"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]} + + +def _responses_turn(gateway: Gateway, group: str, request_input: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/responses", + {"model": group, "input": request_input, "store": False, "include": ["reasoning.encrypted_content"]}, + ) + + +def _turn_one(gateway: Gateway, group: str, deployment_ids: tuple[str, ...]) -> tuple[str, list[JsonValue]]: + # the proxy caches responses, so the turn-one prompt needs a unique marker or a + # stale body carrying another run's encoded model_id would replay instead + response: Final = _responses_turn(gateway, group, f"hello affinity {uuid.uuid4().hex}") + assert response.status_code == 200, response.text + origin: Final = str(response.headers["x-litellm-model-id"]) + assert origin in deployment_ids, f"turn one served by unknown deployment {origin}: {response.text}" + output: Final = object_value(response.json()).get("output") + assert isinstance(output, list), f"turn one returned no output items: {response.text}" + reasoning: Final = next((item for item in output if object_value(item).get("type") == "reasoning"), None) + assert isinstance(reasoning, dict), f"turn one returned no reasoning item: {response.text}" + assert str(object_value(reasoning)["id"]).startswith("encitem_"), ( + f"affinity encoding did not run on turn one: {reasoning}" + ) + assert isinstance(object_value(reasoning).get("encrypted_content"), str), ( + f"turn one reasoning item has no encrypted_content: {reasoning}" + ) + message: Final = next((item for item in output if object_value(item).get("type") == "message"), None) + assert isinstance(message, dict), f"turn one returned no message item: {response.text}" + return origin, [reasoning, message] + + +def _replay(gateway: Gateway, group: str, items: list[JsonValue]) -> httpx.Response: + return _responses_turn( + gateway, + group, + [_user_message("hello affinity"), *items, _user_message("continue the conversation")], + ) + + +def _model_blocked(gateway: Gateway, model_id: str) -> bool: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + entry: Final = next( + (entry for entry in entries if object_value(object_value(entry)["model_info"])["id"] == model_id), + None, + ) + return entry is not None and object_value(object_value(entry)["model_info"]).get("blocked") is True + + +def test_replayed_encrypted_content_serves_from_sibling_when_origin_blocked(gateway: Gateway) -> None: + """Origin excluded from healthy deployments (admin-blocked, no cooldown) must not 503. + + On unfixed code the origin is a routed-group candidate with no (api_base, api_key) + peer, so the affinity check raises a proxy-level 503 instead of stripping the + encrypted reasoning and dispatching to a healthy sibling. + """ + with gateway.scenario() as scenario: + _enable_affinity_check(scenario) + group, deployment_ids = _multi_region_group(scenario) + origin, items = _turn_one(gateway, group, deployment_ids) + + gateway.post("/model/block", {"model_id": origin}) + eventually(lambda: _model_blocked(gateway, origin), lambda blocked: blocked) + + response: Final = _replay(gateway, group, items) + assert response.status_code == 200, response.text + siblings: Final = tuple(model_id for model_id in deployment_ids if model_id != origin) + assert response.headers.get("x-litellm-model-id") in siblings, ( + f"turn two served by {response.headers.get('x-litellm-model-id')}, " + f"expected a sibling of blocked origin {origin}: {response.text}" + ) + _assert_no_encrypted_reasoning_reached_upstream(gateway, group) + + +def test_replayed_encrypted_content_serves_from_sibling_when_origin_deleted(gateway: Gateway) -> None: + """Origin permanently removed: strip-and-dispatch to a sibling, the shipped behavior.""" + with gateway.scenario() as scenario: + _enable_affinity_check(scenario) + group, deployment_ids = _multi_region_group(scenario) + origin, items = _turn_one(gateway, group, deployment_ids) + + gateway.post("/model/delete", {"id": origin}) + eventually( + lambda: gateway.get("/model/info")["data"], + lambda entries: ( + isinstance(entries, list) + and all(object_value(object_value(entry)["model_info"])["id"] != origin for entry in entries) + ), + ) + + response: Final = _replay(gateway, group, items) + assert response.status_code == 200, response.text + siblings: Final = tuple(model_id for model_id in deployment_ids if model_id != origin) + assert response.headers.get("x-litellm-model-id") in siblings, ( + f"turn two served by {response.headers.get('x-litellm-model-id')}, " + f"expected a sibling of deleted origin {origin}: {response.text}" + ) + _assert_no_encrypted_reasoning_reached_upstream(gateway, group) diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 0a3dcba325a..b8c4e5fe084 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -30,6 +30,7 @@ from pydantic import ValidationError import litellm from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck from litellm.types.router import RetryPolicy, UpdateRouterConfig @@ -203,6 +204,63 @@ def test_update_settings_replaces_toggleable_pre_call_checks(): assert any(isinstance(callback, ModelRateLimitingCheck) for callback in (router.optional_callbacks or [])) +def test_update_settings_clears_omitted_encrypted_content_affinity_check(): + router = _build_router() + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + router.update_settings(optional_pre_call_checks=[]) + + assert not any( + isinstance(callback, EncryptedContentAffinityCheck) for callback in (router.optional_callbacks or []) + ) + assert not any(isinstance(callback, EncryptedContentAffinityCheck) for callback in litellm.callbacks) + + +@pytest.mark.asyncio +async def test_update_settings_turning_off_encrypted_content_affinity_stops_flagging_requests(): + router = _build_router() + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + enabled_request: Final = {"litellm_metadata": {}, "input": "hello"} + await router.async_get_available_deployment(model="test-model", request_kwargs=enabled_request) + assert enabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + + router.update_settings(optional_pre_call_checks=[]) + disabled_request: Final = {"litellm_metadata": {}, "input": "hello"} + await router.async_get_available_deployment(model="test-model", request_kwargs=disabled_request) + assert "encrypted_content_affinity_enabled" not in disabled_request["litellm_metadata"] + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + reenabled_request: Final = {"litellm_metadata": {}, "input": "hello"} + await router.async_get_available_deployment(model="test-model", request_kwargs=reenabled_request) + assert reenabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True + + +def test_update_settings_keeps_per_group_encrypted_content_affinity_when_global_toggle_is_omitted(): + router = litellm.Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "sk-fake", + "api_base": "http://localhost:9999", + }, + } + ], + model_group_affinity_config={"test-model": ["encrypted_content_affinity"]}, + ) + + router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"]) + router.update_settings(optional_pre_call_checks=[]) + + affinity_checks: Final = [ + callback for callback in (router.optional_callbacks or []) if isinstance(callback, EncryptedContentAffinityCheck) + ] + assert len(affinity_checks) == 1 + assert affinity_checks[0].enable_global_affinity is False + + @pytest.mark.asyncio async def test_update_settings_preserves_router_budget_limiting_when_omitted(monkeypatch): async def _disable_periodic_sync(*args, **kwargs): diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index aa34fbd6bf7..836049c88a2 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1285,7 +1285,7 @@ def test_boundary_key_rejects_non_dict_like_inputs(): # --------------------------------------------------------------------------- -# Fail-fast when originating deployment is unavailable and no boundary peer +# Degraded dispatch when the originating deployment is unavailable and no boundary peer # --------------------------------------------------------------------------- @@ -1323,14 +1323,13 @@ def _make_router_mock_with_cooldown( @pytest.mark.asyncio -async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_429(): +async def test_affinity_strips_and_dispatches_when_origin_cooled_for_non_429(): """ Originating deployment is in the router config, in cooldown for a non-429 cause (e.g. a 500), and no boundary peer is configured. The check must - surface this as a 503 (transient, but not rate-limit-specific) rather than - dispatching to a non-peer deployment. + degrade: strip the encrypted reasoning and dispatch to the healthy pool + rather than failing the request. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1366,32 +1365,36 @@ async def test_affinity_raises_service_unavailable_when_origin_cooled_for_non_42 } ] request_kwargs = { - "input": [{"id": encoded_id, "type": "reasoning"}], + "input": [ + { + "id": encoded_id, + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + "gAAAAA-blob", "deployment-a-cooled" + ), + } + ], } - with pytest.raises(ServiceUnavailableError) as excinfo: - await check.async_filter_deployments( - model="gpt-5.4", - healthy_deployments=healthy_only_b, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) - # Public error message intentionally omits the originating model_id to - # avoid an authenticated-caller probing oracle. - assert "deployment-a-cooled" not in str(excinfo.value) - assert excinfo.value.status_code == 503 + assert result is healthy_only_b + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in request_kwargs["input"]) @pytest.mark.asyncio -async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_for_429(): +async def test_affinity_strips_and_dispatches_when_origin_cooled_for_429(): """ - Originating deployment is in cooldown specifically because of a 429. - The check must surface this as a 429 RateLimitError with a Retry-After - header derived from the cooldown's remaining window, so OpenAI-compatible - clients respect the backoff instead of giving up on a 503. + Originating deployment is in cooldown specifically because of a 429 and no + boundary peer is configured. The check degrades the same way: strip the + encrypted reasoning and dispatch to the healthy pool rather than surfacing + a rate-limit error to the caller. """ - from litellm.exceptions import RateLimitError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1431,29 +1434,23 @@ async def test_affinity_raises_rate_limit_with_retry_after_when_origin_cooled_fo "input": [{"id": encoded_id, "type": "reasoning"}], } - with pytest.raises(RateLimitError) as excinfo: - await check.async_filter_deployments( - model="gpt-5.4", - healthy_deployments=healthy_only_b, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) - assert "deployment-a-cooled-429" not in str(excinfo.value) - assert excinfo.value.status_code == 429 - retry_after = excinfo.value.response.headers.get("retry-after") - assert retry_after is not None - assert 1 <= int(retry_after) <= 60 + assert result is healthy_only_b @pytest.mark.asyncio -async def test_affinity_raises_service_unavailable_when_origin_filtered_without_cooldown_entry(): +async def test_affinity_strips_and_dispatches_when_origin_filtered_without_cooldown_entry(): """ Originating deployment is configured but absent from healthy_deployments - with no active cooldown entry. Surface as 503 (we cannot prove the cause - was rate-limiting) rather than guessing 429. + with no active cooldown entry and no boundary peer. The check degrades: + strip the encrypted reasoning and dispatch to the healthy pool. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1480,15 +1477,72 @@ async def test_affinity_raises_service_unavailable_when_origin_filtered_without_ "input": [{"id": encoded_id, "type": "reasoning"}], } - with pytest.raises(ServiceUnavailableError) as excinfo: - await check.async_filter_deployments( - model="gpt-5.4", - healthy_deployments=healthy_only_b, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=healthy_only_b, + messages=None, + request_kwargs=request_kwargs, + ) - assert excinfo.value.status_code == 503 + assert result is healthy_only_b + + +@pytest.mark.asyncio +async def test_affinity_serves_sibling_when_candidate_origin_has_no_boundary_peer(): + """ + Regression for the multi-region group where every deployment has a distinct + api_base: the origin is still a routed-group candidate but absent from + healthy_deployments, and no (api_base, api_key) peer exists. The turn must + degrade on a sibling with its encrypted reasoning stripped, not fail. + """ + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + originating = _make_originating_mock("https://region-a.example.com/v1", "shared-key") + mock_router = _make_router_mock_with_cooldown( + originating, cooldown_entries=[], routed_group_model_ids=["region-a", "region-b", "region-c"] + ) + check = EncryptedContentAffinityCheck(router=mock_router) + wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "region-a") + siblings = [ + { + "model_info": {"id": "region-b"}, + "model_name": "gpt-5.4", + "litellm_params": {"api_base": "https://region-b.example.com/v1", "api_key": "shared-key"}, + }, + { + "model_info": {"id": "region-c"}, + "model_name": "gpt-5.4", + "litellm_params": {"api_base": "https://region-c.example.com/v1", "api_key": "shared-key"}, + }, + ] + request_kwargs = { + "litellm_metadata": {}, + "input": [ + {"role": "user", "content": "why is the sky blue?"}, + { + "type": "reasoning", + "id": ResponsesAPIRequestUtils._build_encrypted_item_id("region-a", "rs_test"), + "encrypted_content": wrapped, + "summary": [{"type": "summary_text", "text": "scattering"}], + }, + {"role": "user", "content": "and sunsets?"}, + ], + } + + result = await check.async_filter_deployments( + model="gpt-5.4", + healthy_deployments=siblings, + messages=None, + request_kwargs=request_kwargs, + ) + + assert result is siblings + assert request_kwargs["input"][1] == { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "scattering"}], + } @pytest.mark.asyncio @@ -1981,15 +2035,14 @@ async def test_affinity_strips_encrypted_reasoning_when_routed_to_another_model_ @pytest.mark.asyncio -async def test_affinity_fails_fast_within_the_origins_own_group(): +async def test_affinity_degrades_within_the_origins_own_group(): """ - Negative class for the tier-change discriminator: the routed group IS the - origin's group (a same-group cooldown, not a tier change), so even with a - healthy non-origin sibling that cannot decrypt the content, the request - still fails fast and the encrypted reasoning is left intact rather than - stripped. Preserves the LIT-3051 cooldown contract. + The routed group IS the origin's group (a same-group cooldown), and the + healthy sibling sits on a different encryption boundary, so it cannot + decrypt the replayed reasoning. The check still degrades instead of + failing: the encrypted reasoning is stripped and the request dispatches + to the sibling. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2014,29 +2067,26 @@ async def test_affinity_fails_fast_within_the_origins_own_group(): ] request_kwargs = _cross_group_request_kwargs() - with pytest.raises(ServiceUnavailableError): - await check.async_filter_deployments( - model="gpt-reasoning-tier", - healthy_deployments=sibling_pool, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-reasoning-tier", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert request_kwargs["input"][1].get("encrypted_content") + assert result is sibling_pool + assert not request_kwargs["input"][1].get("encrypted_content") @pytest.mark.asyncio -async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_same_by_id(): +async def test_affinity_strips_when_group_is_spelled_differently_but_same_by_id(): """ - The discriminator must key on deployment-id membership, not on the model-group - name string. Here the origin's configured group is spelled ``openai/gpt-5.4-mini`` - while the routed group is the canonical ``gpt-5.4-mini``: same group, different - spelling. A name compare (``originating.model_name != model``) would read this as - a tier change and strip the reasoning it did not have to. Because the origin's id - is a member of the routed group, this is a same-group cooldown instead: the request - fails fast and the encrypted reasoning is left intact. + The origin's configured group is spelled ``openai/gpt-5.4-mini`` while the + routed group is the canonical ``gpt-5.4-mini``: same group, different + spelling, with the origin a member but currently unavailable. There is no + boundary peer, so the check degrades: strip the encrypted reasoning and + dispatch to the sibling pool. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2071,28 +2121,24 @@ async def test_affinity_does_not_strip_when_group_is_spelled_differently_but_sam ], } - with pytest.raises(ServiceUnavailableError): - await check.async_filter_deployments( - model="gpt-5.4-mini", - healthy_deployments=sibling_pool, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="gpt-5.4-mini", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert request_kwargs["input"][1].get("encrypted_content") + assert result is sibling_pool + assert not request_kwargs["input"][1].get("encrypted_content") @pytest.mark.asyncio -async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes(): +async def test_affinity_strips_for_team_and_pattern_routes(): """ - The exact `model_name` index does not include team-public or pattern routes, so a - same-group cooldown reached only through one of those would be misread as a tier change - and stripped. The check asks the router for the candidate ids it resolves for the route - (`get_candidate_model_ids_for_route`), which covers those paths, rather than the bare - index. Here that set marks the origin as a candidate, so the request fails fast with its - reasoning intact, and the routed group and team are passed through to the router. + Team-public or pattern routes resolve the same routed group as the origin's, + so an unavailable origin there also degrades rather than failing: the + encrypted reasoning is stripped and the request dispatches to the sibling. """ - from litellm.exceptions import ServiceUnavailableError from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -2127,13 +2173,12 @@ async def test_affinity_honors_router_candidate_ids_for_team_and_pattern_routes( ], } - with pytest.raises(ServiceUnavailableError): - await check.async_filter_deployments( - model="team-public-model", - healthy_deployments=sibling_pool, - messages=None, - request_kwargs=request_kwargs, - ) + result = await check.async_filter_deployments( + model="team-public-model", + healthy_deployments=sibling_pool, + messages=None, + request_kwargs=request_kwargs, + ) - assert request_kwargs["input"][1].get("encrypted_content") - mock_router.get_candidate_model_ids_for_route.assert_called_once_with(model="team-public-model", team_id="teamA") + assert result is sibling_pool + assert not request_kwargs["input"][1].get("encrypted_content") From 0d115ad8b50f32443a5785407addf4f0e5cb2693 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:41:02 -0700 Subject: [PATCH 5/6] chore(cost-map): sync gemini priority, flex and video token prices from the Gemini API pricing page (#43091) * chore(cost-map): add gemini priority and flex prices to nano-banana-pro-preview and video token price to 3.1 flash live Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(cost-map): add video token price to bare gemini-3.1-flash-live-preview Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 ++++++ model_prices_and_context_window.json | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c58ea51db9f..333225842cc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -28398,6 +28398,8 @@ "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "gemini", "max_input_tokens": 131072, "max_output_tokens": 32768, @@ -28409,6 +28411,8 @@ "rpm": 1000, "tpm": 4000000, "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -56211,6 +56215,7 @@ "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, "input_cost_per_video_per_second": 3.3333333333333335e-05, + "input_cost_per_video_token": 1e-06, "litellm_provider": "gemini", "max_input_tokens": 131072, "max_output_tokens": 65536, @@ -56403,6 +56408,7 @@ "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, "input_cost_per_video_per_second": 3.3333333333333335e-05, + "input_cost_per_video_token": 1e-06, "litellm_provider": "gemini", "max_input_tokens": 131072, "max_output_tokens": 65536, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c58ea51db9f..333225842cc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -28398,6 +28398,8 @@ "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "gemini", "max_input_tokens": 131072, "max_output_tokens": 32768, @@ -28409,6 +28411,8 @@ "rpm": 1000, "tpm": 4000000, "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -56211,6 +56215,7 @@ "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, "input_cost_per_video_per_second": 3.3333333333333335e-05, + "input_cost_per_video_token": 1e-06, "litellm_provider": "gemini", "max_input_tokens": 131072, "max_output_tokens": 65536, @@ -56403,6 +56408,7 @@ "input_cost_per_image_token": 1e-06, "input_cost_per_token": 7.5e-07, "input_cost_per_video_per_second": 3.3333333333333335e-05, + "input_cost_per_video_token": 1e-06, "litellm_provider": "gemini", "max_input_tokens": 131072, "max_output_tokens": 65536, From 37f52679914154133ae6fcd82f8614125229f0fa Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:59:45 -0700 Subject: [PATCH 6/6] feat(cost-map): add fireworks deepseek-v4p1-flash US-only rows (#43097) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 34 +++++++++++++++++++ model_prices_and_context_window.json | 34 +++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 333225842cc..aa79cadae2a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -60150,6 +60150,23 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { + "cache_read_input_token_cost": 9e-09, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.8e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "deprecation_date": "2026-09-25", @@ -60229,6 +60246,23 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/deepseek-v4p1-flash-us": { + "cache_read_input_token_cost": 9e-09, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.8e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "deprecation_date": "2026-09-25", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 333225842cc..aa79cadae2a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -60150,6 +60150,23 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/accounts/fireworks/routers/deepseek-v4p1-flash-us": { + "cache_read_input_token_cost": 9e-09, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.8e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "deprecation_date": "2026-09-25", @@ -60229,6 +60246,23 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/deepseek-v4p1-flash-us": { + "cache_read_input_token_cost": 9e-09, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 393216, + "max_tokens": 393216, + "mode": "chat", + "output_cost_per_token": 1.8e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, "deprecation_date": "2026-09-25",