From 78927138d288e6a780c83816b05f677b39b9e38e Mon Sep 17 00:00:00 2001
From: michelligabriele <gabriele.michelli@icloud.com>
Date: Thu, 12 Mar 2026 22:30:52 +0100
Subject: [PATCH 01/31] fix(proxy): add team_member_budget_duration to
 NewTeamRequest

NewTeamRequest was missing the team_member_budget_duration field,
causing Pydantic to silently drop the value when creating a team via
POST /team/new. The template budget row was created without
budget_duration or budget_reset_at, so the ResetBudgetJob never
found it and team member spend was never reset.

Add the field to NewTeamRequest and pass it through to
should_create_budget and create_team_member_budget_table in the
new_team handler (matching the existing update_team path which
already works correctly).

Fixes #16057
---
 litellm/proxy/_types.py                       |  1 +
 .../management_endpoints/team_endpoints.py    |  2 +
 .../test_team_endpoints.py                    | 50 +++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 1174740948..9b94a81eb0 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1669,6 +1669,7 @@ class NewTeamRequest(TeamBase):
         int
     ] = None  # allow user to set TPM limit for all team members
     team_member_key_duration: Optional[str] = None  # e.g. "1d", "1w", "1m"
+    team_member_budget_duration: Optional[str] = None  # e.g. "30d", "1mo"
     allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
     enforced_batch_output_expires_after: Optional[dict] = None
     enforced_file_expires_after: Optional[dict] = None
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 9c8e6f7282..18e41793df 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -898,6 +898,7 @@ async def new_team(  # noqa: PLR0915
             team_member_budget=data.team_member_budget,
             team_member_rpm_limit=data.team_member_rpm_limit,
             team_member_tpm_limit=data.team_member_tpm_limit,
+            team_member_budget_duration=data.team_member_budget_duration,
         ):
             data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
                 data=data,
@@ -906,6 +907,7 @@ async def new_team(  # noqa: PLR0915
                 team_member_budget=data.team_member_budget,
                 team_member_rpm_limit=data.team_member_rpm_limit,
                 team_member_tpm_limit=data.team_member_tpm_limit,
+                team_member_budget_duration=data.team_member_budget_duration,
             )
 
         ## ADD TO TEAM TABLE
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index 1aee1d4965..8a033fca96 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -6157,3 +6157,53 @@ async def test_list_team_v1_batches_key_queries():
         assert result[0].keys == [key1, key2]
         assert result[1].team_id == "team-2"
         assert result[1].keys == [key3]
+
+
+def test_new_team_request_accepts_team_member_budget_duration():
+    """Test that NewTeamRequest does not silently drop team_member_budget_duration."""
+    from litellm.proxy._types import NewTeamRequest
+
+    request = NewTeamRequest(
+        team_member_budget=20.0,
+        team_member_budget_duration="30d",
+    )
+    assert request.team_member_budget == 20.0
+    assert request.team_member_budget_duration == "30d"
+
+
+@pytest.mark.asyncio
+async def test_create_team_member_budget_table_with_duration():
+    """Verify that create_team_member_budget_table passes budget_duration
+    through to the new_budget call when team_member_budget_duration is provided."""
+    from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
+    from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
+
+    mock_budget_response = MagicMock(budget_id="budget-abc")
+    mock_admin = UserAPIKeyAuth(
+        user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+
+    data = NewTeamRequest(
+        team_alias="test-team",
+        team_member_budget=20.0,
+        team_member_budget_duration="30d",
+    )
+
+    with patch(
+        "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
+        new_callable=AsyncMock,
+        return_value=mock_budget_response,
+    ) as mock_new_budget:
+        result = await TeamMemberBudgetHandler.create_team_member_budget_table(
+            data=data,
+            new_team_data_json={"metadata": None},
+            user_api_key_dict=mock_admin,
+            team_member_budget=20.0,
+            team_member_budget_duration="30d",
+        )
+
+        mock_new_budget.assert_awaited_once()
+        budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
+        assert budget_request.budget_duration == "30d"
+        assert budget_request.max_budget == 20.0
+        assert result["metadata"]["team_member_budget_id"] == "budget-abc"

From 0b07f628ffb2d4fe8d2c271fe578a9a2796382ec Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 10:30:03 -0700
Subject: [PATCH 02/31] [Test] UI: Add vitest coverage for 10 previously
 untested components

Add unit tests for:
- SimpleToolCallBlock, SimpleMessageBlock, CollapsibleMessage, HistoryTree (log details drawer)
- OnboardingForm (onboarding flow)
- TeamsHeaderTabs, TeamsTable (teams page)
- transform_key_info, filter_helpers (key/team helpers)
- queryKeysFactory (query key generation utility)

47 new tests covering conditional rendering, user interactions, data transformation, and error handling.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../hooks/common/queryKeysFactory.test.ts     |  34 +++++
 .../teams/components/TeamsHeaderTabs.test.tsx |  54 ++++++++
 .../components/TeamsTable/TeamsTable.test.tsx | 129 ++++++++++++++++++
 .../app/onboarding/OnboardingForm.test.tsx    |  95 +++++++++++++
 .../key_team_helpers/filter_helpers.test.ts   |  90 ++++++++++++
 .../transform_key_info.test.ts                |  62 +++++++++
 .../CollapsibleMessage.test.tsx               |  54 ++++++++
 .../LogDetailsDrawer/HistoryTree.test.tsx     |  50 +++++++
 .../SimpleMessageBlock.test.tsx               |  55 ++++++++
 .../SimpleToolCallBlock.test.tsx              |  51 +++++++
 10 files changed, 674 insertions(+)
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx
 create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts
 create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts
 create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx
 create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx
 create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx
 create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts
new file mode 100644
index 0000000000..39afd04409
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts
@@ -0,0 +1,34 @@
+import { describe, it, expect } from "vitest";
+import { createQueryKeys } from "./queryKeysFactory";
+
+describe("createQueryKeys", () => {
+  const keys = createQueryKeys("books");
+
+  it("should return the resource name as the base key", () => {
+    expect(keys.all).toEqual(["books"]);
+  });
+
+  it("should generate a lists key", () => {
+    expect(keys.lists()).toEqual(["books", "list"]);
+  });
+
+  it("should generate a list key with params", () => {
+    expect(keys.list({ page: 1, limit: 10 })).toEqual([
+      "books",
+      "list",
+      { params: { page: 1, limit: 10 } },
+    ]);
+  });
+
+  it("should generate a list key with undefined params when none provided", () => {
+    expect(keys.list()).toEqual(["books", "list", { params: undefined }]);
+  });
+
+  it("should generate a details key", () => {
+    expect(keys.details()).toEqual(["books", "detail"]);
+  });
+
+  it("should generate a detail key for a specific ID", () => {
+    expect(keys.detail("123")).toEqual(["books", "detail", "123"]);
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx
new file mode 100644
index 0000000000..50a7f10f04
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx
@@ -0,0 +1,54 @@
+import { render, screen } from "@testing-library/react";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import TeamsHeaderTabs from "./TeamsHeaderTabs";
+
+vi.mock("@tremor/react", () => ({
+  TabGroup: ({ children, ...props }: any) => <div data-testid="tab-group" {...props}>{children}</div>,
+  TabList: ({ children, ...props }: any) => <div data-testid="tab-list" {...props}>{children}</div>,
+  Tab: ({ children, ...props }: any) => <button {...props}>{children}</button>,
+  TabPanels: ({ children, ...props }: any) => <div data-testid="tab-panels" {...props}>{children}</div>,
+  Text: ({ children, ...props }: any) => <span {...props}>{children}</span>,
+  Icon: ({ onClick, ...props }: any) => <button data-testid="refresh-icon" onClick={onClick} />,
+}));
+
+vi.mock("@heroicons/react/outline", () => ({
+  RefreshIcon: () => <svg data-testid="refresh-svg" />,
+}));
+
+const renderTabs = (props: Partial<Parameters<typeof TeamsHeaderTabs>[0]> = {}) => {
+  const defaults = {
+    lastRefreshed: "",
+    onRefresh: vi.fn(),
+    userRole: "Internal User",
+    children: <div data-testid="panel-content">Panel</div>,
+  };
+  return render(<TeamsHeaderTabs {...defaults} {...props} />);
+};
+
+describe("TeamsHeaderTabs", () => {
+  it("should render 'Your Teams' and 'Available Teams' tabs", () => {
+    renderTabs();
+
+    expect(screen.getByText("Your Teams")).toBeInTheDocument();
+    expect(screen.getByText("Available Teams")).toBeInTheDocument();
+  });
+
+  it("should render 'Default Team Settings' tab when user is Admin", () => {
+    renderTabs({ userRole: "Admin" });
+
+    expect(screen.getByText("Default Team Settings")).toBeInTheDocument();
+  });
+
+  it("should not render 'Default Team Settings' tab for non-admin users", () => {
+    renderTabs({ userRole: "Internal User" });
+
+    expect(screen.queryByText("Default Team Settings")).not.toBeInTheDocument();
+  });
+
+  it("should display last refreshed time when provided", () => {
+    renderTabs({ lastRefreshed: "2024-06-01 12:00:00" });
+
+    expect(screen.getByText("Last Refreshed: 2024-06-01 12:00:00")).toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx
new file mode 100644
index 0000000000..6b072ababb
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx
@@ -0,0 +1,129 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { Team } from "@/components/key_team_helpers/key_list";
+import TeamsTable from "./TeamsTable";
+
+vi.mock("@tremor/react", () => ({
+  Button: React.forwardRef<HTMLButtonElement, any>(({ children, ...props }, ref) =>
+    React.createElement("button", { ...props, ref }, children),
+  ),
+  Icon: ({ onClick, ...props }: any) => <button data-testid={props["data-testid"] || "icon-btn"} onClick={onClick} aria-label={props["aria-label"]} />,
+  Table: ({ children }: any) => <table>{children}</table>,
+  TableHead: ({ children }: any) => <thead>{children}</thead>,
+  TableBody: ({ children }: any) => <tbody>{children}</tbody>,
+  TableRow: ({ children }: any) => <tr>{children}</tr>,
+  TableHeaderCell: ({ children }: any) => <th>{children}</th>,
+  TableCell: ({ children, ...props }: any) => <td {...props}>{children}</td>,
+  Text: ({ children }: any) => <span>{children}</span>,
+}));
+
+vi.mock("antd", () => ({
+  Tooltip: ({ children }: any) => <>{children}</>,
+}));
+
+vi.mock("@heroicons/react/outline", () => ({
+  PencilAltIcon: () => <svg data-testid="pencil-icon" />,
+  TrashIcon: () => <svg data-testid="trash-icon" />,
+}));
+
+vi.mock("@/utils/dataUtils", () => ({
+  formatNumberWithCommas: (val: number, decimals: number) =>
+    val != null ? val.toFixed(decimals) : "N/A",
+}));
+
+vi.mock("@/app/(dashboard)/teams/components/TeamsTable/ModelsCell", () => ({
+  default: ({ team }: any) => <td data-testid="models-cell">{team.models.join(",")}</td>,
+}));
+
+vi.mock("@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell", () => ({
+  default: ({ team }: any) => <td data-testid="role-cell">{team.team_id}</td>,
+}));
+
+const makeTeam = (overrides: Partial<Team> = {}): Team => ({
+  team_id: "team-abc1234",
+  team_alias: "Platform",
+  models: ["gpt-4"],
+  max_budget: 500,
+  budget_duration: null,
+  tpm_limit: null,
+  rpm_limit: null,
+  organization_id: "org-1",
+  created_at: "2024-06-01T00:00:00Z",
+  keys: [],
+  members_with_roles: [],
+  spend: 123.4567,
+  ...overrides,
+});
+
+const defaultPerTeamInfo = {
+  "team-abc1234": {
+    keys: [{ token: "tok-1" } as any, { token: "tok-2" } as any],
+    team_info: {
+      members_with_roles: [{ user_id: "u1", role: "admin" } as any],
+    },
+  },
+};
+
+const renderTable = (overrides: Partial<Parameters<typeof TeamsTable>[0]> = {}) => {
+  const defaults = {
+    teams: [makeTeam()],
+    currentOrg: null,
+    perTeamInfo: defaultPerTeamInfo,
+    userRole: "Admin",
+    userId: "user-1",
+    setSelectedTeamId: vi.fn(),
+    setEditTeam: vi.fn(),
+    onDeleteTeam: vi.fn(),
+  };
+  return render(<TeamsTable {...defaults} {...overrides} />);
+};
+
+describe("TeamsTable", () => {
+  it("should render table headers", () => {
+    renderTable();
+
+    expect(screen.getByText("Team Name")).toBeInTheDocument();
+    expect(screen.getByText("Team ID")).toBeInTheDocument();
+    expect(screen.getByText("Created")).toBeInTheDocument();
+    expect(screen.getByText("Spend (USD)")).toBeInTheDocument();
+    expect(screen.getByText("Budget (USD)")).toBeInTheDocument();
+    expect(screen.getByText("Models")).toBeInTheDocument();
+    expect(screen.getByText("Organization")).toBeInTheDocument();
+    expect(screen.getByText("Your Role")).toBeInTheDocument();
+    expect(screen.getByText("Info")).toBeInTheDocument();
+  });
+
+  it("should render team rows with team data", () => {
+    renderTable();
+
+    expect(screen.getByText("Platform")).toBeInTheDocument();
+    expect(screen.getByText("team-ab...")).toBeInTheDocument();
+    expect(screen.getByText("org-1")).toBeInTheDocument();
+  });
+
+  it("should show edit and delete icons for Admin users", () => {
+    renderTable({ userRole: "Admin" });
+
+    expect(screen.getAllByTestId("icon-btn").length).toBeGreaterThanOrEqual(2);
+  });
+
+  it("should not show edit and delete icons for non-Admin users", () => {
+    renderTable({ userRole: "Internal User" });
+
+    // Only the team ID button should be present, no icon-btn for edit/delete
+    const iconBtns = screen.queryAllByTestId("icon-btn");
+    expect(iconBtns).toHaveLength(0);
+  });
+
+  it("should call setSelectedTeamId when team ID button is clicked", async () => {
+    const user = userEvent.setup();
+    const setSelectedTeamId = vi.fn();
+    renderTable({ setSelectedTeamId });
+
+    await user.click(screen.getByText("team-ab..."));
+
+    expect(setSelectedTeamId).toHaveBeenCalledWith("team-abc1234");
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx
new file mode 100644
index 0000000000..d9e75388d1
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx
@@ -0,0 +1,95 @@
+import { render, screen } from "@testing-library/react";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { OnboardingForm } from "./OnboardingForm";
+
+const mockUseOnboardingCredentials = vi.fn();
+const mockClaimToken = vi.fn();
+
+vi.mock("next/navigation", () => ({
+  useSearchParams: () => new URLSearchParams("invitation_id=inv-123"),
+}));
+
+vi.mock("jwt-decode", () => ({
+  jwtDecode: vi.fn(() => ({
+    user_email: "alice@example.com",
+    user_id: "user-1",
+    key: "access-tok",
+  })),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/onboarding/useOnboarding", () => ({
+  useOnboardingCredentials: (...args: unknown[]) => mockUseOnboardingCredentials(...args),
+  useClaimOnboardingToken: () => ({ mutate: mockClaimToken, isPending: false }),
+}));
+
+vi.mock("@/components/networking", () => ({
+  getProxyBaseUrl: vi.fn(() => ""),
+}));
+
+vi.mock("./OnboardingLoadingView", () => ({
+  OnboardingLoadingView: () => <div data-testid="loading-view">Loading</div>,
+}));
+
+vi.mock("./OnboardingErrorView", () => ({
+  OnboardingErrorView: () => <div data-testid="error-view">Error</div>,
+}));
+
+vi.mock("./OnboardingFormBody", () => ({
+  OnboardingFormBody: ({ variant, userEmail }: { variant: string; userEmail: string }) => (
+    <div data-testid="form-body" data-variant={variant} data-email={userEmail}>
+      Form Body
+    </div>
+  ),
+}));
+
+describe("OnboardingForm", () => {
+  it("should render loading view when credentials are loading", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: undefined,
+      isLoading: true,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("loading-view")).toBeInTheDocument();
+  });
+
+  it("should render error view when credentials fail to load", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: undefined,
+      isLoading: false,
+      isError: true,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("error-view")).toBeInTheDocument();
+  });
+
+  it("should render form body with decoded email when credentials are loaded", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: { token: "fake-jwt-token" },
+      isLoading: false,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("form-body")).toBeInTheDocument();
+    expect(screen.getByTestId("form-body")).toHaveAttribute("data-email", "alice@example.com");
+  });
+
+  it("should pass variant prop to OnboardingFormBody", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: { token: "fake-jwt-token" },
+      isLoading: false,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="reset_password" />);
+
+    expect(screen.getByTestId("form-body")).toHaveAttribute("data-variant", "reset_password");
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts
new file mode 100644
index 0000000000..e3f6a5989f
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts
@@ -0,0 +1,90 @@
+import { describe, expect, it, vi } from "vitest";
+import { fetchTeamFilterOptions } from "./filter_helpers";
+
+const mockKeyListCall = vi.fn();
+
+vi.mock("@/components/networking", () => ({
+  keyListCall: (...args: unknown[]) => mockKeyListCall(...args),
+  teamListCall: vi.fn(),
+  organizationListCall: vi.fn(),
+}));
+
+describe("fetchTeamFilterOptions", () => {
+  it("should return empty arrays when accessToken is null", async () => {
+    const result = await fetchTeamFilterOptions(null, "team-1");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+    expect(mockKeyListCall).not.toHaveBeenCalled();
+  });
+
+  it("should return empty arrays when teamId is empty", async () => {
+    const result = await fetchTeamFilterOptions("tok-123", "");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+    expect(mockKeyListCall).not.toHaveBeenCalled();
+  });
+
+  it("should return sorted key aliases from fetched keys", async () => {
+    mockKeyListCall.mockResolvedValue({
+      keys: [
+        { key_alias: "zeta-key" },
+        { key_alias: "alpha-key" },
+        { key_alias: "mid-key" },
+      ],
+      total_pages: 1,
+    });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.keyAliases).toEqual(["alpha-key", "mid-key", "zeta-key"]);
+  });
+
+  it("should deduplicate organization IDs across pages", async () => {
+    mockKeyListCall
+      .mockResolvedValueOnce({
+        keys: [
+          { organization_id: "org-b" },
+          { organization_id: "org-a" },
+        ],
+        total_pages: 2,
+      })
+      .mockResolvedValueOnce({
+        keys: [
+          { organization_id: "org-a" },
+          { organization_id: "org-c" },
+        ],
+        total_pages: 2,
+      });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.organizationIds).toEqual(["org-a", "org-b", "org-c"]);
+  });
+
+  it("should map user IDs with email addresses", async () => {
+    mockKeyListCall.mockResolvedValue({
+      keys: [
+        { user_id: "u1", user: { user_email: "alice@example.com" } },
+        { user_id: "u2", user: { user_email: "bob@example.com" } },
+      ],
+      total_pages: 1,
+    });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.userIds).toEqual(
+      expect.arrayContaining([
+        { id: "u1", email: "alice@example.com" },
+        { id: "u2", email: "bob@example.com" },
+      ]),
+    );
+  });
+
+  it("should handle API errors gracefully and return empty arrays", async () => {
+    mockKeyListCall.mockRejectedValue(new Error("Network error"));
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts b/ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts
new file mode 100644
index 0000000000..a1139addbf
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts
@@ -0,0 +1,62 @@
+import { describe, it, expect } from "vitest";
+import { transformKeyInfo } from "./transform_key_info";
+
+describe("transformKeyInfo", () => {
+  it("should combine key and info fields into a single object", () => {
+    const apiResponse = {
+      key: "sk-abc123",
+      info: {
+        token_id: "tok_1",
+        key_name: "my-key",
+        spend: 10.5,
+      },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result).toEqual({
+      token: "sk-abc123",
+      token_id: "tok_1",
+      key_name: "my-key",
+      spend: 10.5,
+    });
+  });
+
+  it("should set the token field from the key property", () => {
+    const apiResponse = {
+      key: "sk-xyz789",
+      info: { key_name: "test" },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token).toBe("sk-xyz789");
+  });
+
+  it("should preserve all info fields in the result", () => {
+    const apiResponse = {
+      key: "sk-abc",
+      info: {
+        token_id: "tok_2",
+        key_name: "prod-key",
+        spend: 42,
+        models: ["gpt-4"],
+        team_id: "team-1",
+        metadata: { env: "production" },
+      },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token_id).toBe("tok_2");
+    expect(result.key_name).toBe("prod-key");
+    expect(result.spend).toBe(42);
+    expect(result.models).toEqual(["gpt-4"]);
+    expect(result.team_id).toBe("team-1");
+    expect(result.metadata).toEqual({ env: "production" });
+  });
+
+  it("should handle empty info object", () => {
+    const apiResponse = {
+      key: "sk-empty",
+      info: {},
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token).toBe("sk-empty");
+    expect(Object.keys(result)).toContain("token");
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx
new file mode 100644
index 0000000000..f07d186c04
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx
@@ -0,0 +1,54 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect } from "vitest";
+import { CollapsibleMessage } from "./CollapsibleMessage";
+
+describe("CollapsibleMessage", () => {
+  it("should return null when content is empty", () => {
+    const { container } = render(
+      <CollapsibleMessage label="SYSTEM" content="" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should return null when content is undefined", () => {
+    const { container } = render(<CollapsibleMessage label="SYSTEM" />);
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render the label and char count", () => {
+    render(<CollapsibleMessage label="SYSTEM" content="Hello" />);
+    expect(screen.getByText("SYSTEM")).toBeInTheDocument();
+    expect(screen.getByText("(5 chars)")).toBeInTheDocument();
+  });
+
+  it("should show content when defaultExpanded is true", () => {
+    render(
+      <CollapsibleMessage
+        label="SYSTEM"
+        content="Visible text"
+        defaultExpanded={true}
+      />
+    );
+    expect(screen.getByText("Visible text")).toBeInTheDocument();
+  });
+
+  it("should toggle expanded state when header is clicked", async () => {
+    const user = userEvent.setup();
+    render(
+      <CollapsibleMessage
+        label="SYSTEM"
+        content="Toggle me"
+        defaultExpanded={false}
+      />
+    );
+
+    // Content is rendered in DOM but collapsed by default
+    expect(screen.getByText("Toggle me")).toBeInTheDocument();
+
+    // Click the header to expand - should still show content
+    await user.click(screen.getByText("SYSTEM"));
+    expect(screen.getByText("Toggle me")).toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx
new file mode 100644
index 0000000000..f5eb1fcf8d
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx
@@ -0,0 +1,50 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect } from "vitest";
+import { HistoryTree } from "./HistoryTree";
+import { ParsedMessage } from "./prettyMessagesTypes";
+
+describe("HistoryTree", () => {
+  it("should return null when messages array is empty", () => {
+    const { container } = render(<HistoryTree messages={[]} />);
+    expect(container.innerHTML).toBe("");
+  });
+
+  it('should render message count with plural "messages" for multiple messages', () => {
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there" },
+      { role: "user", content: "How are you?" },
+    ];
+    render(<HistoryTree messages={messages} />);
+    expect(
+      screen.getByText("HISTORY (3 messages)")
+    ).toBeInTheDocument();
+  });
+
+  it('should render message count with singular "message" for one message', () => {
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+    ];
+    render(<HistoryTree messages={messages} />);
+    expect(
+      screen.getByText("HISTORY (1 message)")
+    ).toBeInTheDocument();
+  });
+
+  it("should expand and show messages when header is clicked", async () => {
+    const user = userEvent.setup();
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there" },
+    ];
+    render(<HistoryTree messages={messages} />);
+
+    // Click to expand
+    await user.click(screen.getByText("HISTORY (2 messages)"));
+
+    expect(screen.getByText("Hello")).toBeInTheDocument();
+    expect(screen.getByText("Hi there")).toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx
new file mode 100644
index 0000000000..6483507938
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx
@@ -0,0 +1,55 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { SimpleMessageBlock } from "./SimpleMessageBlock";
+
+describe("SimpleMessageBlock", () => {
+  it("should render the label and content", () => {
+    render(<SimpleMessageBlock label="USER" content="Hello world" />);
+    expect(screen.getByText("USER")).toBeInTheDocument();
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+  });
+
+  it("should return null when content is empty and no tool calls", () => {
+    const { container } = render(
+      <SimpleMessageBlock label="USER" content="" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it('should return null when content is "null" string and no tool calls', () => {
+    const { container } = render(
+      <SimpleMessageBlock label="USER" content="null" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render tool calls when present", () => {
+    render(
+      <SimpleMessageBlock
+        label="ASSISTANT"
+        toolCalls={[
+          { id: "tc1", name: "get_weather", arguments: { city: "Paris" } },
+        ]}
+      />
+    );
+    expect(screen.getByText("ASSISTANT")).toBeInTheDocument();
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+
+  it("should render content and tool calls together", () => {
+    render(
+      <SimpleMessageBlock
+        label="ASSISTANT"
+        content="Let me check the weather."
+        toolCalls={[
+          { id: "tc1", name: "get_weather", arguments: { city: "Paris" } },
+        ]}
+      />
+    );
+    expect(
+      screen.getByText("Let me check the weather.")
+    ).toBeInTheDocument();
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx
new file mode 100644
index 0000000000..33ff403090
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx
@@ -0,0 +1,51 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { SimpleToolCallBlock } from "./SimpleToolCallBlock";
+
+describe("SimpleToolCallBlock", () => {
+  it("should render the tool name", () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+
+  it('should display "function" badge', () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    expect(screen.getByText("function")).toBeInTheDocument();
+  });
+
+  it("should render arguments when present", () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{
+          id: "1",
+          name: "get_weather",
+          arguments: { city: "London", units: "metric" },
+        }}
+      />
+    );
+    expect(screen.getByText("city:")).toBeInTheDocument();
+    expect(screen.getByText('"London"')).toBeInTheDocument();
+    expect(screen.getByText("units:")).toBeInTheDocument();
+    expect(screen.getByText('"metric"')).toBeInTheDocument();
+  });
+
+  it("should not render arguments section when arguments are empty", () => {
+    const { container } = render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    // The tool name and "function" badge should be there, but no key: value pairs
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+    expect(screen.queryByText(/:$/)).not.toBeInTheDocument();
+  });
+});

From 519afa494cd76762e1ea2b24e3961d6477c74cdd Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 12:57:17 -0700
Subject: [PATCH 03/31] [Refactor] UI - API Reference: Migrate to path-based
 routing

Move the API Reference page from query-param routing (?page=api_ref) to
Next.js path-based routing (/ui/api-reference). Add a LEGACY_REDIRECTS
map in the root page.tsx so users with old bookmarks are seamlessly
redirected. Future page migrations only need one new map entry.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../app/(dashboard)/api-reference/page.tsx    | 10 ++-------
 .../app/(dashboard)/components/Sidebar2.tsx   |  2 +-
 .../hooks/proxySettings/useProxySettings.ts   | 20 +++++++++++++++++
 ui/litellm-dashboard/src/app/page.tsx         | 22 +++++++++++++++----
 .../src/components/leftnav.tsx                |  4 ++--
 .../src/components/page_metadata.ts           |  2 +-
 6 files changed, 44 insertions(+), 16 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx
index f9f26c0eac..02bed1adbe 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx
@@ -1,16 +1,10 @@
 "use client";
 
 import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
-import { useState } from "react";
-
-interface ProxySettings {
-  PROXY_BASE_URL: string;
-  PROXY_LOGOUT_URL: string;
-  LITELLM_UI_API_DOC_BASE_URL?: string | null;
-}
+import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
 
 const APIReferencePage = () => {
-  const [proxySettings, setProxySettings] = useState<ProxySettings>({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "" });
+  const proxySettings = useProxySettings();
 
   return <APIReferenceView proxySettings={proxySettings} />;
 };
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx
index 18ab475f22..27a6e6c13b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx
@@ -195,7 +195,7 @@ const menuItems: MenuItemCfg[] = [
     icon: <UserOutlined style={{ fontSize: 18 }} />,
     roles: all_admin_roles,
   },
-  { key: "14", page: "api_ref", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
+  { key: "14", page: "api-reference", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
   {
     key: "16",
     page: "model-hub-table",
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts
new file mode 100644
index 0000000000..62de696b7a
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts
@@ -0,0 +1,20 @@
+import { useState, useEffect } from "react";
+import { fetchProxySettings } from "@/utils/proxyUtils";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+export default function useProxySettings() {
+  const { accessToken } = useAuthorized();
+  const [proxySettings, setProxySettings] = useState({
+    PROXY_BASE_URL: "",
+    PROXY_LOGOUT_URL: "",
+  });
+
+  useEffect(() => {
+    if (!accessToken) return;
+    fetchProxySettings(accessToken).then((settings) => {
+      if (settings) setProxySettings(settings);
+    });
+  }, [accessToken]);
+
+  return proxySettings;
+}
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 5f2921203f..7ee48317e9 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -1,6 +1,5 @@
 "use client";
 
-import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
 import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
 import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView";
 import PlaygroundPage from "@/app/(dashboard)/playground/page";
@@ -48,7 +47,7 @@ import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, stor
 import { formatUserRole, isAdminRole } from "@/utils/roles";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { jwtDecode } from "jwt-decode";
-import { useSearchParams } from "next/navigation";
+import { useRouter, useSearchParams } from "next/navigation";
 import { Suspense, useEffect, useMemo, useRef, useState } from "react";
 import { ConfigProvider, theme } from "antd";
 
@@ -75,6 +74,15 @@ interface ProxySettings {
   LITELLM_UI_API_DOC_BASE_URL?: string | null;
 }
 
+/**
+ * Map of legacy query-param page keys → new path-based route segments.
+ * When a user visits ?page=<key>, they are redirected to /ui/<value>.
+ * Add entries here as pages are migrated from the if/else chain to path-based routes.
+ */
+const LEGACY_REDIRECTS: Record<string, string> = {
+  api_ref: "api-reference",
+};
+
 function CreateKeyPageContent() {
   const [userRole, setUserRole] = useState("");
   const [premiumUser, setPremiumUser] = useState(false);
@@ -90,6 +98,7 @@ function CreateKeyPageContent() {
   });
 
   const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true);
+  const router = useRouter();
   const searchParams = useSearchParams()!;
   const [modelData, setModelData] = useState<any>({ data: [] });
   const [token, setToken] = useState<string | null>(null);
@@ -431,6 +440,13 @@ function CreateKeyPageContent() {
     return <LoadingScreen />;
   }
 
+  // Redirect legacy query-param pages to their new path-based routes
+  if (page in LEGACY_REDIRECTS) {
+    const base = (proxyBaseUrl || "") + "/ui";
+    router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
+    return <LoadingScreen />;
+  }
+
   return (
     <Suspense fallback={<LoadingScreen />}>
       <ConfigProvider theme={{
@@ -536,8 +552,6 @@ function CreateKeyPageContent() {
                     <AdminPanel
                       proxySettings={proxySettings}
                     />
-                  ) : page == "api_ref" ? (
-                    <APIReferenceView proxySettings={proxySettings} />
                   ) : page == "logging-and-alerts" ? (
                     <Settings userID={userID} userRole={userRole} accessToken={accessToken} premiumUser={premiumUser} />
                   ) : page == "budgets" ? (
diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx
index d01fc06bc0..30b3a9653f 100644
--- a/ui/litellm-dashboard/src/components/leftnav.tsx
+++ b/ui/litellm-dashboard/src/components/leftnav.tsx
@@ -231,8 +231,8 @@ const menuGroups: MenuGroup[] = [
     groupLabel: "DEVELOPER TOOLS",
     items: [
       {
-        key: "api_ref",
-        page: "api_ref",
+        key: "api-reference",
+        page: "api-reference",
         label: "API Reference",
         icon: <ApiOutlined />,
       },
diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts
index a910373d66..fdfc321f38 100644
--- a/ui/litellm-dashboard/src/components/page_metadata.ts
+++ b/ui/litellm-dashboard/src/components/page_metadata.ts
@@ -24,7 +24,7 @@ export const pageDescriptions: Record<string, string> = {
   projects: "Manage projects within teams",
   "access-groups": "Manage access groups for role-based permissions",
   budgets: "Set and monitor spending budgets",
-  api_ref: "Browse API documentation and endpoints",
+  "api-reference": "Browse API documentation and endpoints",
   "model-hub-table": "Explore available AI models and providers",
   "learning-resources": "Access tutorials and documentation",
   caching: "Configure response caching settings",

From 895951744ddbf62ef31a78b86d953d7f6eb6cefb Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 13:36:21 -0700
Subject: [PATCH 04/31] Fix leftnav navigation regression and useProxySettings
 initial state

The leftnav was updated to emit page="api-reference" but only "api_ref"
was in LEGACY_REDIRECTS, causing clicks to fall through to the default
Usage page. Add "api-reference" entry to the redirect map. Also include
LITELLM_UI_API_DOC_BASE_URL in the hook's initial state to avoid a
brief flash of incorrect base URL.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts  | 1 +
 ui/litellm-dashboard/src/app/page.tsx                            | 1 +
 2 files changed, 2 insertions(+)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts
index 62de696b7a..d4fb307385 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts
@@ -7,6 +7,7 @@ export default function useProxySettings() {
   const [proxySettings, setProxySettings] = useState({
     PROXY_BASE_URL: "",
     PROXY_LOGOUT_URL: "",
+    LITELLM_UI_API_DOC_BASE_URL: null as string | null,
   });
 
   useEffect(() => {
diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 7ee48317e9..5f545b635e 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -81,6 +81,7 @@ interface ProxySettings {
  */
 const LEGACY_REDIRECTS: Record<string, string> = {
   api_ref: "api-reference",
+  "api-reference": "api-reference",
 };
 
 function CreateKeyPageContent() {

From e9e5ed989c29103ab2af839add0a1e822553820b Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 13:59:45 -0700
Subject: [PATCH 05/31] Move legacy redirect from render phase to useEffect

router.replace was called directly during render, which is unsafe in
React 18 concurrent mode. Move it into a useEffect and use a computed
flag (isLegacyRedirect) to show LoadingScreen while redirecting.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ui/litellm-dashboard/src/app/page.tsx | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 5f545b635e..3aac962135 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -253,6 +253,15 @@ function CreateKeyPageContent() {
     }
   }, [redirectToLogin]);
 
+  // Redirect legacy query-param pages to their new path-based routes
+  const isLegacyRedirect = page in LEGACY_REDIRECTS;
+  useEffect(() => {
+    if (isLegacyRedirect) {
+      const base = (proxyBaseUrl || "") + "/ui";
+      router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
+    }
+  }, [isLegacyRedirect, page, router]);
+
   // Check for a stored return URL after successful authentication
   // This handles the case where user comes back from SSO and we need to redirect to the original URL
   useEffect(() => {
@@ -437,14 +446,7 @@ function CreateKeyPageContent() {
     setShowClaudeCodePrompt(true);
   };
 
-  if (authLoading || redirectToLogin) {
-    return <LoadingScreen />;
-  }
-
-  // Redirect legacy query-param pages to their new path-based routes
-  if (page in LEGACY_REDIRECTS) {
-    const base = (proxyBaseUrl || "") + "/ui";
-    router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
+  if (authLoading || redirectToLogin || isLegacyRedirect) {
     return <LoadingScreen />;
   }
 

From 2e70c23307d85c7430bdf0523204863865797d79 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 14:05:56 -0700
Subject: [PATCH 06/31] Gate legacy redirect on authLoading to ensure
 proxyBaseUrl is resolved

The redirect useEffect fires before getUiConfig() completes, so
proxyBaseUrl is always "" on first render. Gate on !authLoading so
the redirect only fires after config is fetched, matching the pattern
used by the login redirect.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ui/litellm-dashboard/src/app/page.tsx | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 3aac962135..03431428ad 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -256,11 +256,11 @@ function CreateKeyPageContent() {
   // Redirect legacy query-param pages to their new path-based routes
   const isLegacyRedirect = page in LEGACY_REDIRECTS;
   useEffect(() => {
-    if (isLegacyRedirect) {
+    if (!authLoading && isLegacyRedirect) {
       const base = (proxyBaseUrl || "") + "/ui";
       router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
     }
-  }, [isLegacyRedirect, page, router]);
+  }, [authLoading, isLegacyRedirect, page, router]);
 
   // Check for a stored return URL after successful authentication
   // This handles the case where user comes back from SSO and we need to redirect to the original URL

From f3b17686160893b888fdf27f15290b2dbcbe2fe0 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 17:40:50 -0700
Subject: [PATCH 07/31] feat: add _assert_can_manage_team_mcp_server auth
 helper with tests

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/proxy/_types.py                       |   3 +-
 .../management_endpoints/common_utils.py      |  12 ++
 .../mcp_management_endpoints.py               |  69 ++++++-
 .../test_mcp_manager_role.py                  | 187 ++++++++++++++++++
 4 files changed, 269 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 9e86680e35..0c2cb8c827 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1608,8 +1608,9 @@ class Member(MemberBase):
     role: Literal[
         "admin",
         "user",
+        "mcp_server_manager",
     ] = Field(
-        description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member"
+        description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member, 'mcp_server_manager' can manage MCP servers for the team"
     )
 
 
diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py
index efc42d3355..092b22cee6 100644
--- a/litellm/proxy/management_endpoints/common_utils.py
+++ b/litellm/proxy/management_endpoints/common_utils.py
@@ -41,6 +41,18 @@ def _is_user_team_admin(
     return False
 
 
+def _is_user_team_mcp_manager(
+    user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
+) -> bool:
+    for member in team_obj.members_with_roles:
+        if (
+            member.user_id is not None and member.user_id == user_api_key_dict.user_id
+        ) and member.role == "mcp_server_manager":
+            return True
+
+    return False
+
+
 async def _is_user_org_admin_for_team(
     user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
 ) -> bool:
diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index 3e5b729cea..b94bd87596 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -130,11 +130,78 @@ if MCP_AVAILABLE:
     )
     from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
     from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
-    from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
+    from litellm.proxy.auth.auth_checks import get_team_object
+    from litellm.proxy.management_endpoints.common_utils import (
+        _is_user_team_mcp_manager,
+        _user_has_admin_view,
+    )
+    from litellm.proxy.management_helpers.object_permission_utils import (
+        _get_team_allowed_mcp_servers,
+        handle_update_object_permission_common,
+    )
     from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
     from litellm.types.mcp import MCPCredentials
     from litellm.types.mcp_server.mcp_server_manager import MCPServer
 
+    async def _assert_can_manage_team_mcp_server(
+        user_api_key_dict: UserAPIKeyAuth,
+        team_id: Optional[str] = None,
+        server_id: Optional[str] = None,
+    ) -> str:
+        """
+        Verify that the caller is an MCP server manager for a team and (for edit/delete)
+        that the target server belongs to that team.
+
+        Returns the team_id the caller is managing.
+        Raises HTTPException(400) if no team_id can be determined.
+        Raises HTTPException(403) if the caller is not an MCP manager or server not in team.
+        """
+        from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
+
+        resolved_team_id = team_id or user_api_key_dict.team_id
+        if not resolved_team_id:
+            raise HTTPException(
+                status_code=400,
+                detail={"error": "team_id is required for MCP server manager operations."},
+            )
+
+        if (
+            team_id
+            and user_api_key_dict.team_id
+            and team_id != user_api_key_dict.team_id
+        ):
+            raise HTTPException(
+                status_code=403,
+                detail={"error": "team_id does not match the API key's team."},
+            )
+
+        team_obj = await get_team_object(
+            team_id=resolved_team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            check_db_only=False,
+        )
+
+        if not _is_user_team_mcp_manager(user_api_key_dict, team_obj):
+            raise HTTPException(
+                status_code=403,
+                detail={
+                    "error": f"User does not have mcp_server_manager role in team {resolved_team_id}."
+                },
+            )
+
+        if server_id is not None:
+            team_server_ids = await _get_team_allowed_mcp_servers(team_obj)
+            if server_id not in team_server_ids:
+                raise HTTPException(
+                    status_code=403,
+                    detail={
+                        "error": f"MCP server {server_id} is not assigned to team {resolved_team_id}."
+                    },
+                )
+
+        return resolved_team_id
+
     @dataclass
     class _TemporaryMCPServerEntry:
         server: MCPServer
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
new file mode 100644
index 0000000000..20cb337cd1
--- /dev/null
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -0,0 +1,187 @@
+import pytest
+
+from litellm.proxy._types import (
+    LiteLLM_TeamTable,
+    LitellmUserRoles,
+    Member,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.management_endpoints.common_utils import (
+    _is_user_team_mcp_manager,
+)
+
+
+class TestIsUserTeamMcpManager:
+    def test_mcp_server_manager_role_returns_true(self):
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+        )
+        team = LiteLLM_TeamTable(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager")
+            ],
+        )
+        assert _is_user_team_mcp_manager(user_auth, team) is True
+
+    def test_regular_user_role_returns_false(self):
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+        )
+        team = LiteLLM_TeamTable(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="user")],
+        )
+        assert _is_user_team_mcp_manager(user_auth, team) is False
+
+    def test_admin_role_returns_false(self):
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+        )
+        team = LiteLLM_TeamTable(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="admin")],
+        )
+        assert _is_user_team_mcp_manager(user_auth, team) is False
+
+    def test_user_not_in_team_returns_false(self):
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user2",
+            api_key="sk-test",
+        )
+        team = LiteLLM_TeamTable(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager")
+            ],
+        )
+        assert _is_user_team_mcp_manager(user_auth, team) is False
+
+
+from unittest.mock import AsyncMock, MagicMock, patch
+from litellm.proxy._types import LiteLLM_TeamTableCachedObj
+
+
+@pytest.mark.asyncio
+class TestAssertCanManageTeamMcpServer:
+    async def test_mcp_manager_with_team_id_succeeds(self):
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            _assert_can_manage_team_mcp_server,
+        )
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test",
+        )
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="mcp_server_manager")],
+        )
+        with patch(
+            "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
+            AsyncMock(return_value=mock_team),
+        ):
+            result = await _assert_can_manage_team_mcp_server(
+                user_api_key_dict=user_auth, team_id="team1"
+            )
+            assert result == "team1"
+
+    async def test_regular_user_gets_403(self):
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            _assert_can_manage_team_mcp_server,
+        )
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test",
+        )
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="user")],
+        )
+        with patch(
+            "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
+            AsyncMock(return_value=mock_team),
+        ):
+            with pytest.raises(Exception) as exc_info:
+                await _assert_can_manage_team_mcp_server(
+                    user_api_key_dict=user_auth, team_id="team1"
+                )
+            assert exc_info.value.status_code == 403
+
+    async def test_admin_gets_403(self):
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            _assert_can_manage_team_mcp_server,
+        )
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test",
+        )
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="admin")],
+        )
+        with patch(
+            "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
+            AsyncMock(return_value=mock_team),
+        ):
+            with pytest.raises(Exception) as exc_info:
+                await _assert_can_manage_team_mcp_server(
+                    user_api_key_dict=user_auth, team_id="team1"
+                )
+            assert exc_info.value.status_code == 403
+
+    async def test_mcp_manager_server_in_team_succeeds(self):
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            _assert_can_manage_team_mcp_server,
+        )
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1",
+        )
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="mcp_server_manager")],
+        )
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
+                AsyncMock(return_value=mock_team),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers",
+                AsyncMock(return_value={"server1", "server2"}),
+            ),
+        ):
+            result = await _assert_can_manage_team_mcp_server(
+                user_api_key_dict=user_auth, server_id="server1"
+            )
+            assert result == "team1"
+
+    async def test_mcp_manager_server_not_in_team_gets_403(self):
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            _assert_can_manage_team_mcp_server,
+        )
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1",
+        )
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[Member(user_id="user1", role="mcp_server_manager")],
+        )
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
+                AsyncMock(return_value=mock_team),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers",
+                AsyncMock(return_value={"server2", "server3"}),
+            ),
+        ):
+            with pytest.raises(Exception) as exc_info:
+                await _assert_can_manage_team_mcp_server(
+                    user_api_key_dict=user_auth, server_id="server1"
+                )
+            assert exc_info.value.status_code == 403

From cc7e27e094d0d0c3f1e3b31a73675d1cb096823c Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 17:55:06 -0700
Subject: [PATCH 08/31] feat: update MCP CRUD endpoints to support
 mcp_server_manager role

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../mcp_management_endpoints.py               | 89 +++++++++++++------
 .../test_mcp_manager_role.py                  | 84 +++++++++++++++++
 2 files changed, 147 insertions(+), 26 deletions(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index b94bd87596..3400f90f48 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -1273,16 +1273,26 @@ if MCP_AVAILABLE:
         # Validate and normalize payload fields
         validate_and_normalize_mcp_server_payload(payload)
 
-        # AuthZ - restrict only proxy admins to create mcp servers
-        if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
-            raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail={
-                    "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN."
-                },
+        # AuthZ - proxy admins or team MCP managers can create MCP servers
+        is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role
+        manager_team_id: Optional[str] = None
+
+        if not is_proxy_admin:
+            # Check if the user is an MCP manager for the specified team
+            if payload.team_id is None:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={
+                        "error": "team_id is required when creating MCP servers as a team MCP manager."
+                    },
+                )
+            manager_team_id = await _assert_can_manage_team_mcp_server(
+                user_api_key_dict=user_api_key_dict,
+                team_id=payload.team_id,
             )
-        elif payload.server_id is not None:
-            # fail if the mcp server with id already exists
+
+        # Fail if the MCP server with this id already exists
+        if payload.server_id is not None:
             mcp_server = await get_mcp_server(prisma_client, payload.server_id)
             if mcp_server is not None:
                 raise HTTPException(
@@ -1291,7 +1301,8 @@ if MCP_AVAILABLE:
                         "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
                     },
                 )
-        elif (
+
+        if (
             SpecialMCPServerName.all_team_servers == payload.server_id
             or SpecialMCPServerName.all_proxy_servers == payload.server_id
         ):
@@ -1322,6 +1333,40 @@ if MCP_AVAILABLE:
 
             # Ensure registry is up to date by reloading from database
             await global_mcp_server_manager.reload_servers_from_database()
+
+            # If created by an MCP manager, auto-assign the server to their team
+            if manager_team_id is not None:
+                from litellm.proxy.proxy_server import user_api_key_cache
+
+                team_obj = await get_team_object(
+                    team_id=manager_team_id,
+                    prisma_client=prisma_client,
+                    user_api_key_cache=user_api_key_cache,
+                    check_db_only=False,
+                )
+                existing_permission_id = getattr(
+                    team_obj, "object_permission_id", None
+                )
+
+                # Read existing mcp_servers list and append the new server
+                existing_mcp_servers: list = []
+                if team_obj.object_permission is not None:
+                    existing_mcp_servers = (
+                        team_obj.object_permission.mcp_servers or []
+                    )
+                updated_mcp_servers = list(
+                    set(existing_mcp_servers + [new_mcp_server.server_id])
+                )
+
+                await handle_update_object_permission_common(
+                    data_json={
+                        "object_permission": {
+                            "mcp_servers": updated_mcp_servers,
+                        }
+                    },
+                    existing_object_permission_id=existing_permission_id,
+                    prisma_client=prisma_client,
+                )
         except Exception as e:
             verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}")
             raise HTTPException(
@@ -1540,15 +1585,11 @@ if MCP_AVAILABLE:
             "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
         )
 
-        # Authz - restrict only admins to delete mcp servers
+        # Authz - proxy admins or team MCP managers can delete MCP servers
         if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
-            raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail={
-                    "error": "Call not allowed to delete MCP server. User is not a proxy admin. route={}".format(
-                        "DELETE /v1/mcp/server"
-                    )
-                },
+            await _assert_can_manage_team_mcp_server(
+                user_api_key_dict=user_api_key_dict,
+                server_id=server_id,
             )
 
         # try to delete the mcp server
@@ -1865,15 +1906,11 @@ if MCP_AVAILABLE:
         # Validate and normalize payload fields
         validate_and_normalize_mcp_server_payload(payload)
 
-        # Authz - restrict only admins to delete mcp servers
+        # Authz - proxy admins or team MCP managers can update MCP servers
         if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
-            raise HTTPException(
-                status_code=status.HTTP_403_FORBIDDEN,
-                detail={
-                    "error": "Call not allowed to update MCP server. User is not a proxy admin. route={}".format(
-                        "PUT /v1/mcp/server"
-                    )
-                },
+            await _assert_can_manage_team_mcp_server(
+                user_api_key_dict=user_api_key_dict,
+                server_id=payload.server_id,
             )
 
         # try to update the mcp server
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
index 20cb337cd1..860b255f9c 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -185,3 +185,87 @@ class TestAssertCanManageTeamMcpServer:
                     user_api_key_dict=user_auth, server_id="server1"
                 )
             assert exc_info.value.status_code == 403
+
+
+@pytest.mark.asyncio
+class TestCreateMcpServerAsManager:
+    async def test_create_auto_assigns_to_team(self):
+        """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable."""
+        from litellm.proxy._types import NewMCPServerRequest
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            add_mcp_server,
+        )
+
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+            team_id="team1",
+        )
+
+        payload = NewMCPServerRequest(
+            server_name="test-server",
+            url="https://example.com/mcp",
+            team_id="team1",
+        )
+
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager"),
+            ],
+            object_permission_id="perm1",
+        )
+        mock_team.object_permission = MagicMock(mcp_servers=["existing_server"])
+
+        created_server = MagicMock()
+        created_server.server_id = "new_server_id"
+        created_server.credentials = None
+
+        mock_handle_update = AsyncMock()
+
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+                AsyncMock(return_value="team1"),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
+                AsyncMock(return_value=None),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server",
+                AsyncMock(return_value=created_server),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+                MagicMock(
+                    add_server=AsyncMock(),
+                    reload_servers_from_database=AsyncMock(),
+                ),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
+                AsyncMock(return_value=mock_team),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
+                mock_handle_update,
+            ),
+        ):
+            await add_mcp_server(payload=payload, user_api_key_dict=user_auth)
+
+            # Verify handle_update_object_permission_common was called with merged server list
+            mock_handle_update.assert_called_once()
+            call_kwargs = mock_handle_update.call_args.kwargs
+            mcp_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"]
+            assert "existing_server" in mcp_servers
+            assert "new_server_id" in mcp_servers
+            assert call_kwargs["existing_object_permission_id"] == "perm1"

From 971cc2a9af40fa14b7a50efb428aab9098cc13ec Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 17:56:08 -0700
Subject: [PATCH 09/31] fix: add team_id field to NewMCPServerRequest

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/proxy/_types.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 0c2cb8c827..67ae31bea8 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1144,6 +1144,10 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
         None,
         description="Server-managed: set by the endpoint; caller values are overridden.",
     )
+    team_id: Optional[str] = Field(
+        None,
+        description="Team ID to assign the MCP server to. Required for team MCP managers.",
+    )
 
     @model_validator(mode="before")
     @classmethod

From f3c6915d617d3e8c7ea44547a2ce42accf51e87f Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 17:56:33 -0700
Subject: [PATCH 10/31] feat: add useChatHistory hook with tests (extracted
 from ChatUI)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../playground/chat_ui/useChatHistory.test.ts | 506 ++++++++++++++++++
 .../playground/chat_ui/useChatHistory.ts      | 402 ++++++++++++++
 ui/litellm-dashboard/tests/setupTests.ts      |   5 +
 3 files changed, 913 insertions(+)
 create mode 100644 ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
 create mode 100644 ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts

diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
new file mode 100644
index 0000000000..312d22c28e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
@@ -0,0 +1,506 @@
+import { renderHook, act } from "@testing-library/react";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { useChatHistory } from "./useChatHistory";
+
+describe("useChatHistory", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+  });
+
+  describe("updateTextUI", () => {
+    it("should create a new assistant message when chat is empty", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "Hello", model: "gpt-4" },
+      ]);
+    });
+
+    it("should append to the last assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", " world");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "Hello world", model: "gpt-4" },
+      ]);
+    });
+
+    it("should not overwrite model on subsequent chunks", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", " world", "gpt-3.5");
+      });
+
+      expect(result.current.chatHistory[0].model).toBe("gpt-4");
+    });
+
+    it("should create a new message when role changes", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("user", "Hi");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[0].role).toBe("user");
+      expect(result.current.chatHistory[1].role).toBe("assistant");
+    });
+
+    it("should not append to image messages", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateImageUI("http://img.png", "dall-e");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "description", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+    });
+
+    it("should not append to audio messages", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("http://audio.mp3", "tts-1");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "text", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+    });
+  });
+
+  describe("updateReasoningContent", () => {
+    it("should add reasoning content to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Answer", "gpt-4");
+      });
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory[0].reasoningContent).toBe("thinking...");
+    });
+
+    it("should append reasoning content across chunks", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "", "gpt-4");
+      });
+      act(() => {
+        result.current.updateReasoningContent("step 1");
+      });
+      act(() => {
+        result.current.updateReasoningContent(" step 2");
+      });
+
+      expect(result.current.chatHistory[0].reasoningContent).toBe("step 1 step 2");
+    });
+
+    it("should create assistant message with reasoning when last message is user", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.setChatHistory([{ role: "user", content: "question" }]);
+      });
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[1]).toEqual({
+        role: "assistant",
+        content: "",
+        reasoningContent: "thinking...",
+      });
+    });
+
+    it("should not update when chat is empty", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(0);
+    });
+  });
+
+  describe("updateTimingData", () => {
+    it("should add timeToFirstToken to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTimingData(150);
+      });
+
+      expect(result.current.chatHistory[0].timeToFirstToken).toBe(150);
+    });
+
+    it("should create assistant message when last is user", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.setChatHistory([{ role: "user", content: "hi" }]);
+      });
+      act(() => {
+        result.current.updateTimingData(200);
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[1].timeToFirstToken).toBe(200);
+    });
+  });
+
+  describe("updateUsageData", () => {
+    it("should add usage data to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
+      act(() => {
+        result.current.updateUsageData(usage);
+      });
+
+      expect(result.current.chatHistory[0].usage).toEqual(usage);
+    });
+
+    it("should add toolName when provided", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
+      act(() => {
+        result.current.updateUsageData(usage, "search_tool");
+      });
+
+      expect(result.current.chatHistory[0].toolName).toBe("search_tool");
+    });
+  });
+
+  describe("updateTotalLatency", () => {
+    it("should add totalLatency to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTotalLatency(500);
+      });
+
+      expect(result.current.chatHistory[0].totalLatency).toBe(500);
+    });
+  });
+
+  describe("updateA2AMetadata", () => {
+    it("should add A2A metadata to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const metadata = { taskId: "task-1", contextId: "ctx-1" };
+      act(() => {
+        result.current.updateA2AMetadata(metadata);
+      });
+
+      expect(result.current.chatHistory[0].a2aMetadata).toEqual(metadata);
+    });
+  });
+
+  describe("updateSearchResults", () => {
+    it("should add search results to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const searchResults = [{ object: "search", search_query: "test", data: [] }];
+      act(() => {
+        result.current.updateSearchResults(searchResults);
+      });
+
+      expect(result.current.chatHistory[0].searchResults).toEqual(searchResults);
+    });
+  });
+
+  describe("updateImageUI", () => {
+    it("should add image message to history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateImageUI("http://img.png", "dall-e-3");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "http://img.png", model: "dall-e-3", isImage: true },
+      ]);
+    });
+  });
+
+  describe("updateEmbeddingsUI", () => {
+    it("should add truncated embeddings message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateEmbeddingsUI("[0.1, 0.2, 0.3]", "text-embedding-ada");
+      });
+
+      expect(result.current.chatHistory[0].isEmbeddings).toBe(true);
+      expect(result.current.chatHistory[0].model).toBe("text-embedding-ada");
+    });
+  });
+
+  describe("updateAudioUI", () => {
+    it("should add audio message to history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("http://audio.mp3", "tts-1");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "http://audio.mp3", model: "tts-1", isAudio: true },
+      ]);
+    });
+  });
+
+  describe("updateChatImageUI", () => {
+    it("should add image to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Here is the image", "gpt-4");
+      });
+      act(() => {
+        result.current.updateChatImageUI("http://img.png", "gpt-4");
+      });
+
+      expect(result.current.chatHistory[0].image).toEqual({
+        url: "http://img.png",
+        detail: "auto",
+      });
+    });
+
+    it("should create new assistant message with image when no assistant message exists", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateChatImageUI("http://img.png", "gpt-4");
+      });
+
+      expect(result.current.chatHistory[0]).toEqual({
+        role: "assistant",
+        content: "",
+        model: "gpt-4",
+        image: { url: "http://img.png", detail: "auto" },
+      });
+    });
+  });
+
+  describe("handleMCPEvent", () => {
+    it("should add MCP event", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(1);
+    });
+
+    it("should deduplicate events by item_id and type", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      const event = { type: "tool_call", item_id: "1" };
+      act(() => {
+        result.current.handleMCPEvent(event);
+      });
+      act(() => {
+        result.current.handleMCPEvent(event);
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(1);
+    });
+
+    it("should allow events without item_id (no dedup)", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call" });
+      });
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call" });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(2);
+    });
+
+    it("should allow events with same item_id/type but different sequence_number", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 1 });
+      });
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 2 });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(2);
+    });
+  });
+
+  describe("clearMCPEvents", () => {
+    it("should clear MCP events without affecting chat history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+      act(() => {
+        result.current.clearMCPEvents();
+      });
+
+      expect(result.current.mcpEvents).toEqual([]);
+      expect(result.current.chatHistory).toHaveLength(1);
+    });
+  });
+
+  describe("clearChatHistory", () => {
+    it("should clear all state", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(result.current.chatHistory).toEqual([]);
+      expect(result.current.mcpEvents).toEqual([]);
+      expect(result.current.messageTraceId).toBeNull();
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("should revoke audio object URLs when clearing", () => {
+      const revokeSpy = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {});
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("blob:http://localhost/audio-1", "tts-1");
+      });
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(revokeSpy).toHaveBeenCalledWith("blob:http://localhost/audio-1");
+      revokeSpy.mockRestore();
+    });
+
+    it("should clear sessionStorage when not simplified", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      sessionStorage.setItem("chatHistory", "[]");
+      sessionStorage.setItem("messageTraceId", "trace-1");
+      sessionStorage.setItem("responsesSessionId", "resp-1");
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+    });
+
+    it("should NOT clear sessionStorage when simplified", () => {
+      sessionStorage.setItem("chatHistory", '[{"role":"user","content":"hi"}]');
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // simplified mode should not touch sessionStorage
+      expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
+    });
+  });
+
+  describe("session management", () => {
+    it("handleResponseId should set responsesSessionId when useApiSessionManagement is true", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+
+      expect(result.current.responsesSessionId).toBe("resp-123");
+    });
+
+    it("handleResponseId should NOT set responsesSessionId when useApiSessionManagement is false", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleToggleSessionManagement(false);
+      });
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("handleToggleSessionManagement should clear session when switching to UI mode", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+      act(() => {
+        result.current.handleToggleSessionManagement(false);
+      });
+
+      expect(result.current.useApiSessionManagement).toBe(false);
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
new file mode 100644
index 0000000000..8830b8321c
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
@@ -0,0 +1,402 @@
+import { useState, useEffect } from "react";
+import { MessageType, A2ATaskMetadata } from "./types";
+import { TokenUsage } from "./ResponseMetrics";
+import { MCPEvent } from "../../mcp_tools/types";
+import { truncateString } from "../../../utils/textUtils";
+
+export interface UseChatHistoryReturn {
+  // State
+  chatHistory: MessageType[];
+  setChatHistory: React.Dispatch<React.SetStateAction<MessageType[]>>;
+  mcpEvents: MCPEvent[];
+  setMCPEvents: React.Dispatch<React.SetStateAction<MCPEvent[]>>;
+  messageTraceId: string | null;
+  setMessageTraceId: React.Dispatch<React.SetStateAction<string | null>>;
+  responsesSessionId: string | null;
+  setResponsesSessionId: React.Dispatch<React.SetStateAction<string | null>>;
+  useApiSessionManagement: boolean;
+  setUseApiSessionManagement: React.Dispatch<React.SetStateAction<boolean>>;
+
+  // Actions
+  updateTextUI: (role: string, chunk: string, model?: string) => void;
+  updateReasoningContent: (chunk: string) => void;
+  updateTimingData: (timeToFirstToken: number) => void;
+  updateUsageData: (usage: TokenUsage, toolName?: string) => void;
+  updateA2AMetadata: (a2aMetadata: A2ATaskMetadata) => void;
+  updateTotalLatency: (totalLatency: number) => void;
+  updateSearchResults: (searchResults: any[]) => void;
+  handleResponseId: (responseId: string) => void;
+  handleToggleSessionManagement: (useApi: boolean) => void;
+  handleMCPEvent: (event: MCPEvent) => void;
+  updateImageUI: (imageUrl: string, model: string) => void;
+  updateEmbeddingsUI: (embeddings: string, model?: string) => void;
+  updateAudioUI: (audioUrl: string, model: string) => void;
+  updateChatImageUI: (imageUrl: string, model?: string) => void;
+  clearChatHistory: () => void;
+  clearMCPEvents: () => void;
+}
+
+export function useChatHistory({ simplified }: { simplified: boolean }): UseChatHistoryReturn {
+  const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
+    if (simplified) return [];
+    try {
+      const saved = sessionStorage.getItem("chatHistory");
+      return saved ? JSON.parse(saved) : [];
+    } catch (error) {
+      console.error("Error parsing chatHistory from sessionStorage", error);
+      return [];
+    }
+  });
+
+  const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
+
+  const [messageTraceId, setMessageTraceId] = useState<string | null>(
+    () => sessionStorage.getItem("messageTraceId") || null,
+  );
+
+  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
+    () => sessionStorage.getItem("responsesSessionId") || null,
+  );
+
+  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
+    const saved = sessionStorage.getItem("useApiSessionManagement");
+    return saved ? JSON.parse(saved) : true; // Default to API session management
+  });
+
+  // Debounced chatHistory persistence
+  useEffect(() => {
+    if (simplified) return; // Do not persist chat history in simplified (embedded) mode
+    const handler = setTimeout(() => {
+      sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
+    }, 500); // Debounce by 500ms
+
+    return () => {
+      clearTimeout(handler);
+    };
+  }, [chatHistory, simplified]);
+
+  // messageTraceId/responsesSessionId/useApiSessionManagement persistence
+  useEffect(() => {
+    if (messageTraceId) {
+      sessionStorage.setItem("messageTraceId", messageTraceId);
+    } else {
+      sessionStorage.removeItem("messageTraceId");
+    }
+    if (responsesSessionId) {
+      sessionStorage.setItem("responsesSessionId", responsesSessionId);
+    } else {
+      sessionStorage.removeItem("responsesSessionId");
+    }
+    sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
+  }, [messageTraceId, responsesSessionId, useApiSessionManagement]);
+
+  const updateTextUI = (role: string, chunk: string, model?: string) => {
+    console.log("updateTextUI called with:", role, chunk, model);
+    setChatHistory((prev) => {
+      const last = prev[prev.length - 1];
+      // if the last message is already from this same role, append
+      if (last && last.role === role && !last.isImage && !last.isAudio) {
+        // build a new object, but only set `model` if it wasn't there already
+        const updated: MessageType = {
+          ...last,
+          content: last.content + chunk,
+          model: last.model ?? model, // ← only use the passed‐in model on the first chunk
+        };
+        return [...prev.slice(0, -1), updated];
+      } else {
+        // otherwise start a brand new assistant bubble
+        return [
+          ...prev,
+          {
+            role,
+            content: chunk,
+            model, // model set exactly once here
+          },
+        ];
+      }
+    });
+  };
+
+  const updateReasoningContent = (chunk: string) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            reasoningContent: (lastMessage.reasoningContent || "") + chunk,
+          },
+        ];
+      } else {
+        // If there's no assistant message yet, we'll create one with empty content
+        // but with reasoning content
+        if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
+          return [
+            ...prevHistory,
+            {
+              role: "assistant",
+              content: "",
+              reasoningContent: chunk,
+            },
+          ];
+        }
+
+        return prevHistory;
+      }
+    });
+  };
+
+  const updateTimingData = (timeToFirstToken: number) => {
+    console.log("updateTimingData called with:", timeToFirstToken);
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+      console.log("Current last message:", lastMessage);
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        console.log("Updating assistant message with timeToFirstToken:", timeToFirstToken);
+        const updatedHistory = [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            timeToFirstToken,
+          },
+        ];
+        console.log("Updated chat history:", updatedHistory);
+        return updatedHistory;
+      }
+      // If the last message is a user message and no assistant message exists yet,
+      // create a new assistant message with empty content
+      else if (lastMessage && lastMessage.role === "user") {
+        console.log("Creating new assistant message with timeToFirstToken:", timeToFirstToken);
+        return [
+          ...prevHistory,
+          {
+            role: "assistant",
+            content: "",
+            timeToFirstToken,
+          },
+        ];
+      }
+
+      console.log("No appropriate message found to update timing");
+      return prevHistory;
+    });
+  };
+
+  const updateUsageData = (usage: TokenUsage, toolName?: string) => {
+    console.log("Received usage data:", usage);
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        console.log("Updating message with usage data:", usage);
+        const updatedMessage = {
+          ...lastMessage,
+          usage,
+          toolName,
+        };
+        console.log("Updated message:", updatedMessage);
+
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
+    console.log("Received A2A metadata:", a2aMetadata);
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          a2aMetadata,
+        };
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateTotalLatency = (totalLatency: number) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            totalLatency,
+          },
+        ];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateSearchResults = (searchResults: any[]) => {
+    console.log("Received search results:", searchResults);
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        console.log("Updating message with search results");
+        const updatedMessage = {
+          ...lastMessage,
+          searchResults,
+        };
+
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const handleResponseId = (responseId: string) => {
+    console.log("Received response ID for session management:", responseId);
+    if (useApiSessionManagement) {
+      setResponsesSessionId(responseId);
+    }
+  };
+
+  const handleToggleSessionManagement = (useApi: boolean) => {
+    setUseApiSessionManagement(useApi);
+    if (!useApi) {
+      // Clear API session when switching to UI mode
+      setResponsesSessionId(null);
+    }
+  };
+
+  const handleMCPEvent = (event: MCPEvent) => {
+    console.log("ChatUI: Received MCP event:", event);
+    setMCPEvents((prev) => {
+      // Check if this is a duplicate event (same item_id and type)
+      // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
+      const isDuplicate = event.item_id
+        ? prev.some(
+            (existingEvent) =>
+              existingEvent.item_id === event.item_id &&
+              existingEvent.type === event.type &&
+              (existingEvent.sequence_number === event.sequence_number ||
+                (existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
+          )
+        : false;
+
+      if (isDuplicate) {
+        console.log("ChatUI: Duplicate MCP event, skipping");
+        return prev;
+      }
+
+      const newEvents = [...prev, event];
+      console.log("ChatUI: Updated MCP events:", newEvents);
+      return newEvents;
+    });
+  };
+
+  const updateImageUI = (imageUrl: string, model: string) => {
+    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
+  };
+
+  const updateEmbeddingsUI = (embeddings: string, model?: string) => {
+    setChatHistory((prevHistory) => [
+      ...prevHistory,
+      { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
+    ]);
+  };
+
+  const updateAudioUI = (audioUrl: string, model: string) => {
+    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
+  };
+
+  const updateChatImageUI = (imageUrl: string, model?: string) => {
+    setChatHistory((prev) => {
+      const last = prev[prev.length - 1];
+      // If the last message is from assistant and has content, add image to it
+      if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
+        const updated = {
+          ...last,
+          image: {
+            url: imageUrl,
+            detail: "auto",
+          },
+          model: last.model ?? model,
+        };
+        return [...prev.slice(0, -1), updated];
+      } else {
+        // Otherwise create a new assistant message with just the image
+        return [
+          ...prev,
+          {
+            role: "assistant",
+            content: "",
+            model,
+            image: {
+              url: imageUrl,
+              detail: "auto",
+            },
+          },
+        ];
+      }
+    });
+  };
+
+  const clearChatHistory = () => {
+    // Clean up audio object URLs before clearing history
+    chatHistory.forEach((message) => {
+      if (message.isAudio && typeof message.content === "string") {
+        URL.revokeObjectURL(message.content);
+      }
+    });
+
+    setChatHistory([]);
+    setMessageTraceId(null);
+    setResponsesSessionId(null); // Clear responses session ID
+    setMCPEvents([]); // Clear MCP events
+    if (!simplified) {
+      sessionStorage.removeItem("chatHistory");
+      sessionStorage.removeItem("messageTraceId");
+      sessionStorage.removeItem("responsesSessionId");
+    }
+  };
+
+  const clearMCPEvents = () => {
+    setMCPEvents([]);
+  };
+
+  return {
+    chatHistory,
+    setChatHistory,
+    mcpEvents,
+    setMCPEvents,
+    messageTraceId,
+    setMessageTraceId,
+    responsesSessionId,
+    setResponsesSessionId,
+    useApiSessionManagement,
+    setUseApiSessionManagement,
+    updateTextUI,
+    updateReasoningContent,
+    updateTimingData,
+    updateUsageData,
+    updateA2AMetadata,
+    updateTotalLatency,
+    updateSearchResults,
+    handleResponseId,
+    handleToggleSessionManagement,
+    handleMCPEvent,
+    updateImageUI,
+    updateEmbeddingsUI,
+    updateAudioUI,
+    updateChatImageUI,
+    clearChatHistory,
+    clearMCPEvents,
+  };
+}
diff --git a/ui/litellm-dashboard/tests/setupTests.ts b/ui/litellm-dashboard/tests/setupTests.ts
index e8de9f7248..5f85a80e58 100644
--- a/ui/litellm-dashboard/tests/setupTests.ts
+++ b/ui/litellm-dashboard/tests/setupTests.ts
@@ -100,6 +100,11 @@ if (!document.getAnimations) {
   document.getAnimations = () => [];
 }
 
+// Stub URL.revokeObjectURL so vi.spyOn can intercept it in tests
+if (!URL.revokeObjectURL) {
+  URL.revokeObjectURL = () => {};
+}
+
 // Mock ResizeObserver for components that use it (e.g., Tremor UI components)
 // This prevents "ResizeObserver is not defined" errors in JSDOM
 global.ResizeObserver = class ResizeObserver {

From e88425b88113d6f7f9a67043c4265be6049876ed Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 18:02:50 -0700
Subject: [PATCH 11/31] refactor: wire ChatUI to use useChatHistory hook

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../components/playground/chat_ui/ChatUI.tsx  | 362 ++----------------
 .../playground/chat_ui/useChatHistory.ts      |   2 +-
 2 files changed, 36 insertions(+), 328 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
index bc0d52cf58..88f3c85c21 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
@@ -28,7 +28,6 @@ import ReactMarkdown from "react-markdown";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
 import { v4 as uuidv4 } from "uuid";
-import { truncateString } from "../../../utils/textUtils";
 import GuardrailSelector from "../../guardrails/GuardrailSelector";
 import PolicySelector from "../../policies/PolicySelector";
 import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/MCPToolArgumentsForm";
@@ -62,7 +61,6 @@ import { generateCodeSnippet } from "./CodeSnippets";
 import EndpointSelector from "./EndpointSelector";
 import FilePreviewCard from "./FilePreviewCard";
 import MCPEventsDisplay from "./MCPEventsDisplay";
-import type { MCPEvent } from "../../mcp_tools/types";
 import { EndpointType, getEndpointType } from "./mode_endpoint_mapping";
 import ReasoningContent from "./ReasoningContent";
 import ResponseMetrics, { TokenUsage } from "./ResponseMetrics";
@@ -74,6 +72,7 @@ import SessionManagement from "./SessionManagement";
 import RealtimePlayground from "./RealtimePlayground";
 import { A2ATaskMetadata, MessageType } from "./types";
 import { useCodeInterpreter } from "./useCodeInterpreter";
+import { useChatHistory } from "./useChatHistory";
 
 const { TextArea } = Input;
 const { Dragger } = Upload;
@@ -134,6 +133,34 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return {};
     }
   });
+  const {
+    chatHistory,
+    setChatHistory,
+    mcpEvents,
+    setMCPEvents,
+    messageTraceId,
+    setMessageTraceId,
+    responsesSessionId,
+    setResponsesSessionId,
+    useApiSessionManagement,
+    setUseApiSessionManagement,
+    updateTextUI,
+    updateReasoningContent,
+    updateTimingData,
+    updateUsageData,
+    updateA2AMetadata,
+    updateTotalLatency,
+    updateSearchResults,
+    handleResponseId,
+    handleToggleSessionManagement,
+    handleMCPEvent,
+    updateImageUI,
+    updateEmbeddingsUI,
+    updateAudioUI,
+    updateChatImageUI,
+    clearChatHistory: clearChatHistoryHook,
+    clearMCPEvents,
+  } = useChatHistory({ simplified });
   const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
     const saved = sessionStorage.getItem("apiKeySource");
     if (saved) {
@@ -150,16 +177,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     () => sessionStorage.getItem("customProxyBaseUrl") || "",
   );
   const [inputMessage, setInputMessage] = useState("");
-  const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
-    if (simplified) return [];
-    try {
-      const saved = sessionStorage.getItem("chatHistory");
-      return saved ? JSON.parse(saved) : [];
-    } catch (error) {
-      console.error("Error parsing chatHistory from sessionStorage", error);
-      return [];
-    }
-  });
   const [selectedModel, setSelectedModel] = useState<string | undefined>(simplified ? fixedModel : undefined);
   const [showCustomModelInput, setShowCustomModelInput] = useState<boolean>(false);
   const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
@@ -217,16 +234,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return [];
     }
   });
-  const [messageTraceId, setMessageTraceId] = useState<string | null>(
-    () => sessionStorage.getItem("messageTraceId") || null,
-  );
-  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
-    () => sessionStorage.getItem("responsesSessionId") || null,
-  );
-  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
-    const saved = sessionStorage.getItem("useApiSessionManagement");
-    return saved ? JSON.parse(saved) : true; // Default to API session management
-  });
   const [uploadedImages, setUploadedImages] = useState<File[]>([]);
   const [imagePreviewUrls, setImagePreviewUrls] = useState<string[]>([]);
   const [responsesUploadedImage, setResponsesUploadedImage] = useState<File | null>(null);
@@ -237,7 +244,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const [isGetCodeModalVisible, setIsGetCodeModalVisible] = useState(false);
   const [generatedCode, setGeneratedCode] = useState("");
   const [selectedSdk, setSelectedSdk] = useState<"openai" | "azure">("openai");
-  const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
   const [temperature, setTemperature] = useState<number>(1.0);
   const [maxTokens, setMaxTokens] = useState<number>(2048);
   const [useAdvancedParams, setUseAdvancedParams] = useState<boolean>(false);
@@ -331,17 +337,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     proxySettings,
   ]);
 
-  useEffect(() => {
-    if (simplified) return; // Do not persist chat history in simplified (embedded) mode
-    const handler = setTimeout(() => {
-      sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
-    }, 500); // Debounce by 500ms
-
-    return () => {
-      clearTimeout(handler);
-    };
-  }, [chatHistory, simplified]);
-
   useEffect(() => {
     sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource));
     sessionStorage.setItem("apiKey", apiKey);
@@ -362,17 +357,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
         sessionStorage.removeItem("selectedModel");
       }
     }
-    if (messageTraceId) {
-      sessionStorage.setItem("messageTraceId", messageTraceId);
-    } else {
-      sessionStorage.removeItem("messageTraceId");
-    }
-    if (responsesSessionId) {
-      sessionStorage.setItem("responsesSessionId", responsesSessionId);
-    } else {
-      sessionStorage.removeItem("responsesSessionId");
-    }
-    sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
     // Note: codeInterpreterEnabled and selectedContainerId are persisted by useCodeInterpreter hook
   }, [
     simplified,
@@ -384,9 +368,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     selectedVectorStores,
     selectedGuardrails,
     selectedPolicies,
-    messageTraceId,
-    responsesSessionId,
-    useApiSessionManagement,
     selectedMCPServers,
     mcpServerToolRestrictions,
     selectedVoice,
@@ -478,264 +459,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
   }, [chatHistory]);
 
-  const updateTextUI = (role: string, chunk: string, model?: string) => {
-    console.log("updateTextUI called with:", role, chunk, model);
-    setChatHistory((prev) => {
-      const last = prev[prev.length - 1];
-      // if the last message is already from this same role, append
-      if (last && last.role === role && !last.isImage && !last.isAudio) {
-        // build a new object, but only set `model` if it wasn't there already
-        const updated: MessageType = {
-          ...last,
-          content: last.content + chunk,
-          model: last.model ?? model, // ← only use the passed‐in model on the first chunk
-        };
-        return [...prev.slice(0, -1), updated];
-      } else {
-        // otherwise start a brand new assistant bubble
-        return [
-          ...prev,
-          {
-            role,
-            content: chunk,
-            model, // model set exactly once here
-          },
-        ];
-      }
-    });
-  };
-
-  const updateReasoningContent = (chunk: string) => {
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
-        return [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            reasoningContent: (lastMessage.reasoningContent || "") + chunk,
-          },
-        ];
-      } else {
-        // If there's no assistant message yet, we'll create one with empty content
-        // but with reasoning content
-        if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
-          return [
-            ...prevHistory,
-            {
-              role: "assistant",
-              content: "",
-              reasoningContent: chunk,
-            },
-          ];
-        }
-
-        return prevHistory;
-      }
-    });
-  };
-
-  const updateTimingData = (timeToFirstToken: number) => {
-    console.log("updateTimingData called with:", timeToFirstToken);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-      console.log("Current last message:", lastMessage);
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating assistant message with timeToFirstToken:", timeToFirstToken);
-        const updatedHistory = [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            timeToFirstToken,
-          },
-        ];
-        console.log("Updated chat history:", updatedHistory);
-        return updatedHistory;
-      }
-      // If the last message is a user message and no assistant message exists yet,
-      // create a new assistant message with empty content
-      else if (lastMessage && lastMessage.role === "user") {
-        console.log("Creating new assistant message with timeToFirstToken:", timeToFirstToken);
-        return [
-          ...prevHistory,
-          {
-            role: "assistant",
-            content: "",
-            timeToFirstToken,
-          },
-        ];
-      }
-
-      console.log("No appropriate message found to update timing");
-      return prevHistory;
-    });
-  };
-
-  const updateUsageData = (usage: TokenUsage, toolName?: string) => {
-    console.log("Received usage data:", usage);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with usage data:", usage);
-        const updatedMessage = {
-          ...lastMessage,
-          usage,
-          toolName,
-        };
-        console.log("Updated message:", updatedMessage);
-
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
-    console.log("Received A2A metadata:", a2aMetadata);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        const updatedMessage = {
-          ...lastMessage,
-          a2aMetadata,
-        };
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateTotalLatency = (totalLatency: number) => {
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        return [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            totalLatency,
-          },
-        ];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateSearchResults = (searchResults: any[]) => {
-    console.log("Received search results:", searchResults);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with search results");
-        const updatedMessage = {
-          ...lastMessage,
-          searchResults,
-        };
-
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const handleResponseId = (responseId: string) => {
-    console.log("Received response ID for session management:", responseId);
-    if (useApiSessionManagement) {
-      setResponsesSessionId(responseId);
-    }
-  };
-
-  const handleToggleSessionManagement = (useApi: boolean) => {
-    setUseApiSessionManagement(useApi);
-    if (!useApi) {
-      // Clear API session when switching to UI mode
-      setResponsesSessionId(null);
-    }
-  };
-
-  const handleMCPEvent = (event: MCPEvent) => {
-    console.log("ChatUI: Received MCP event:", event);
-    setMCPEvents((prev) => {
-      // Check if this is a duplicate event (same item_id and type)
-      // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
-      const isDuplicate = event.item_id
-        ? prev.some(
-            (existingEvent) =>
-              existingEvent.item_id === event.item_id &&
-              existingEvent.type === event.type &&
-              (existingEvent.sequence_number === event.sequence_number ||
-                (existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
-          )
-        : false;
-
-      if (isDuplicate) {
-        console.log("ChatUI: Duplicate MCP event, skipping");
-        return prev;
-      }
-
-      const newEvents = [...prev, event];
-      console.log("ChatUI: Updated MCP events:", newEvents);
-      return newEvents;
-    });
-  };
-
-  const updateImageUI = (imageUrl: string, model: string) => {
-    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
-  };
-
-  const updateEmbeddingsUI = (embeddings: string, model?: string) => {
-    setChatHistory((prevHistory) => [
-      ...prevHistory,
-      { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
-    ]);
-  };
-
-  const updateAudioUI = (audioUrl: string, model: string) => {
-    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
-  };
-
-  const updateChatImageUI = (imageUrl: string, model?: string) => {
-    setChatHistory((prev) => {
-      const last = prev[prev.length - 1];
-      // If the last message is from assistant and has content, add image to it
-      if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
-        const updated = {
-          ...last,
-          image: {
-            url: imageUrl,
-            detail: "auto",
-          },
-          model: last.model ?? model,
-        };
-        return [...prev.slice(0, -1), updated];
-      } else {
-        // Otherwise create a new assistant message with just the image
-        return [
-          ...prev,
-          {
-            role: "assistant",
-            content: "",
-            model,
-            image: {
-              url: imageUrl,
-              detail: "auto",
-            },
-          },
-        ];
-      }
-    });
-  };
-
   const handleKeyDown = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
     if (event.key === "Enter" && !event.shiftKey) {
       event.preventDefault(); // Prevent default to avoid newline
@@ -966,7 +689,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
 
     setChatHistory([...chatHistory, displayMessage]);
-    setMCPEvents([]); // Clear previous MCP events for new conversation turn
+    clearMCPEvents(); // Clear previous MCP events for new conversation turn
     codeInterpreter.clearResult(); // Clear previous code interpreter results
     setIsLoading(true);
 
@@ -1222,26 +945,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
   };
 
   const clearChatHistory = () => {
-    // Clean up audio object URLs before clearing history
-    chatHistory.forEach((message) => {
-      if (message.isAudio && typeof message.content === "string") {
-        URL.revokeObjectURL(message.content);
-      }
-    });
-
-    setChatHistory([]);
-    setMessageTraceId(null);
-    setResponsesSessionId(null); // Clear responses session ID
-    setMCPEvents([]); // Clear MCP events
-    handleRemoveAllImages(); // Clear any uploaded images for image edits
-    handleRemoveResponsesImage(); // Clear any uploaded images for responses
-    handleRemoveChatImage(); // Clear any uploaded images for chat completions
-    handleRemoveAudio(); // Clear any uploaded audio for transcription
-    if (!simplified) {
-      sessionStorage.removeItem("chatHistory");
-      sessionStorage.removeItem("messageTraceId");
-      sessionStorage.removeItem("responsesSessionId");
-    }
+    clearChatHistoryHook();
+    handleRemoveAllImages();
+    handleRemoveResponsesImage();
+    handleRemoveChatImage();
+    handleRemoveAudio();
     NotificationsManager.success("Chat history cleared.");
   };
 
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
index 8830b8321c..52575ef026 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
@@ -1,4 +1,4 @@
-import { useState, useEffect } from "react";
+import React, { useState, useEffect } from "react";
 import { MessageType, A2ATaskMetadata } from "./types";
 import { TokenUsage } from "./ResponseMetrics";
 import { MCPEvent } from "../../mcp_tools/types";

From 954bfcbd421429996db0f71d45908e9ee843203f Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 20:46:51 -0700
Subject: [PATCH 12/31] chore: remove verbose console.log calls from
 useChatHistory hook

Removes 17 console.log statements that fired on every streaming chunk
in updateTextUI, updateTimingData, updateUsageData, and other hot-path
functions. The console.error for sessionStorage parse failures is kept.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../playground/chat_ui/useChatHistory.ts      | 23 ++-----------------
 1 file changed, 2 insertions(+), 21 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
index 52575ef026..36e16c1701 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
@@ -91,7 +91,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   }, [messageTraceId, responsesSessionId, useApiSessionManagement]);
 
   const updateTextUI = (role: string, chunk: string, model?: string) => {
-    console.log("updateTextUI called with:", role, chunk, model);
     setChatHistory((prev) => {
       const last = prev[prev.length - 1];
       // if the last message is already from this same role, append
@@ -149,27 +148,21 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   };
 
   const updateTimingData = (timeToFirstToken: number) => {
-    console.log("updateTimingData called with:", timeToFirstToken);
     setChatHistory((prevHistory) => {
       const lastMessage = prevHistory[prevHistory.length - 1];
-      console.log("Current last message:", lastMessage);
 
       if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating assistant message with timeToFirstToken:", timeToFirstToken);
-        const updatedHistory = [
+        return [
           ...prevHistory.slice(0, prevHistory.length - 1),
           {
             ...lastMessage,
             timeToFirstToken,
           },
         ];
-        console.log("Updated chat history:", updatedHistory);
-        return updatedHistory;
       }
       // If the last message is a user message and no assistant message exists yet,
       // create a new assistant message with empty content
       else if (lastMessage && lastMessage.role === "user") {
-        console.log("Creating new assistant message with timeToFirstToken:", timeToFirstToken);
         return [
           ...prevHistory,
           {
@@ -180,24 +173,20 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
         ];
       }
 
-      console.log("No appropriate message found to update timing");
       return prevHistory;
     });
   };
 
   const updateUsageData = (usage: TokenUsage, toolName?: string) => {
-    console.log("Received usage data:", usage);
     setChatHistory((prevHistory) => {
       const lastMessage = prevHistory[prevHistory.length - 1];
 
       if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with usage data:", usage);
         const updatedMessage = {
           ...lastMessage,
           usage,
           toolName,
         };
-        console.log("Updated message:", updatedMessage);
 
         return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
       }
@@ -207,7 +196,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   };
 
   const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
-    console.log("Received A2A metadata:", a2aMetadata);
     setChatHistory((prevHistory) => {
       const lastMessage = prevHistory[prevHistory.length - 1];
 
@@ -242,12 +230,10 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   };
 
   const updateSearchResults = (searchResults: any[]) => {
-    console.log("Received search results:", searchResults);
     setChatHistory((prevHistory) => {
       const lastMessage = prevHistory[prevHistory.length - 1];
 
       if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with search results");
         const updatedMessage = {
           ...lastMessage,
           searchResults,
@@ -261,7 +247,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   };
 
   const handleResponseId = (responseId: string) => {
-    console.log("Received response ID for session management:", responseId);
     if (useApiSessionManagement) {
       setResponsesSessionId(responseId);
     }
@@ -276,7 +261,6 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   };
 
   const handleMCPEvent = (event: MCPEvent) => {
-    console.log("ChatUI: Received MCP event:", event);
     setMCPEvents((prev) => {
       // Check if this is a duplicate event (same item_id and type)
       // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
@@ -291,13 +275,10 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
         : false;
 
       if (isDuplicate) {
-        console.log("ChatUI: Duplicate MCP event, skipping");
         return prev;
       }
 
-      const newEvents = [...prev, event];
-      console.log("ChatUI: Updated MCP events:", newEvents);
-      return newEvents;
+      return [...prev, event];
     });
   };
 

From f34fe4758a88cae5124348f7863a24e93febfc1a Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 21:19:51 -0700
Subject: [PATCH 13/31] fix: stale closure, simplified session isolation,
 debounce-race in useChatHistory

- clearChatHistory: use functional setChatHistory updater so blob URL
  revocation operates on the latest snapshot, not a stale closure capture.
- Simplified mode: skip sessionStorage hydration and persistence for
  messageTraceId, responsesSessionId, and useApiSessionManagement so
  embedded widgets don't cross-contaminate the full playground session.
- Debounce race: skip re-writing empty chatHistory to sessionStorage
  after clearChatHistory already removed the key.
- Added 5 new tests covering these fixes (39 total).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../playground/chat_ui/useChatHistory.test.ts | 85 +++++++++++++++++++
 .../playground/chat_ui/useChatHistory.ts      | 27 ++++--
 2 files changed, 103 insertions(+), 9 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
index 312d22c28e..c2067fc698 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts
@@ -436,6 +436,7 @@ describe("useChatHistory", () => {
     });
 
     it("should clear sessionStorage when not simplified", () => {
+      vi.useFakeTimers();
       const { result } = renderHook(() => useChatHistory({ simplified: false }));
 
       sessionStorage.setItem("chatHistory", "[]");
@@ -446,9 +447,16 @@ describe("useChatHistory", () => {
         result.current.clearChatHistory();
       });
 
+      // Advance past the 500ms debounce to verify it does not re-write the key
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+
       expect(sessionStorage.getItem("chatHistory")).toBeNull();
       expect(sessionStorage.getItem("messageTraceId")).toBeNull();
       expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+
+      vi.useRealTimers();
     });
 
     it("should NOT clear sessionStorage when simplified", () => {
@@ -463,6 +471,83 @@ describe("useChatHistory", () => {
       // simplified mode should not touch sessionStorage
       expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
     });
+
+    it("should not re-write chatHistory to sessionStorage after clear via debounce", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      // Add a message so the debounce has something to persist
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      // Let the debounce fire so the message is persisted
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+      expect(sessionStorage.getItem("chatHistory")).not.toBeNull();
+
+      // Now clear
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // Advance past the debounce — the key should stay removed
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+
+      vi.useRealTimers();
+    });
+  });
+
+  describe("simplified mode session isolation", () => {
+    it("should not hydrate messageTraceId from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("messageTraceId", "trace-from-playground");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      expect(result.current.messageTraceId).toBeNull();
+    });
+
+    it("should not hydrate responsesSessionId from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("responsesSessionId", "resp-from-playground");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("should not hydrate useApiSessionManagement from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("useApiSessionManagement", "false");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      // Should get the default (true), not the stored value
+      expect(result.current.useApiSessionManagement).toBe(true);
+    });
+
+    it("should not persist session state to sessionStorage in simplified mode", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      act(() => {
+        result.current.setMessageTraceId("trace-embedded");
+        result.current.setResponsesSessionId("resp-embedded");
+      });
+
+      // Flush effects
+      act(() => {
+        vi.advanceTimersByTime(0);
+      });
+
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+
+      vi.useRealTimers();
+    });
   });
 
   describe("session management", () => {
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
index 36e16c1701..8e38191259 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts
@@ -51,14 +51,15 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
 
   const [messageTraceId, setMessageTraceId] = useState<string | null>(
-    () => sessionStorage.getItem("messageTraceId") || null,
+    () => (simplified ? null : sessionStorage.getItem("messageTraceId") || null),
   );
 
   const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
-    () => sessionStorage.getItem("responsesSessionId") || null,
+    () => (simplified ? null : sessionStorage.getItem("responsesSessionId") || null),
   );
 
   const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
+    if (simplified) return true;
     const saved = sessionStorage.getItem("useApiSessionManagement");
     return saved ? JSON.parse(saved) : true; // Default to API session management
   });
@@ -66,6 +67,9 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   // Debounced chatHistory persistence
   useEffect(() => {
     if (simplified) return; // Do not persist chat history in simplified (embedded) mode
+    // When chatHistory is empty (e.g. after clearChatHistory removed the key),
+    // don't re-write an empty array back into sessionStorage.
+    if (chatHistory.length === 0) return;
     const handler = setTimeout(() => {
       sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
     }, 500); // Debounce by 500ms
@@ -77,6 +81,7 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
 
   // messageTraceId/responsesSessionId/useApiSessionManagement persistence
   useEffect(() => {
+    if (simplified) return;
     if (messageTraceId) {
       sessionStorage.setItem("messageTraceId", messageTraceId);
     } else {
@@ -88,7 +93,7 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
       sessionStorage.removeItem("responsesSessionId");
     }
     sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
-  }, [messageTraceId, responsesSessionId, useApiSessionManagement]);
+  }, [messageTraceId, responsesSessionId, useApiSessionManagement, simplified]);
 
   const updateTextUI = (role: string, chunk: string, model?: string) => {
     setChatHistory((prev) => {
@@ -330,14 +335,18 @@ export function useChatHistory({ simplified }: { simplified: boolean }): UseChat
   };
 
   const clearChatHistory = () => {
-    // Clean up audio object URLs before clearing history
-    chatHistory.forEach((message) => {
-      if (message.isAudio && typeof message.content === "string") {
-        URL.revokeObjectURL(message.content);
-      }
+    // Use functional updater to get the latest snapshot — avoids stale-closure
+    // bugs where audio messages added between the last render and the click
+    // would leak their blob URLs.
+    setChatHistory((prev) => {
+      prev.forEach((message) => {
+        if (message.isAudio && typeof message.content === "string") {
+          URL.revokeObjectURL(message.content);
+        }
+      });
+      return [];
     });
 
-    setChatHistory([]);
     setMessageTraceId(null);
     setResponsesSessionId(null); // Clear responses session ID
     setMCPEvents([]); // Clear MCP events

From 883611aa918981e70e1dbe122efabe8a32b3af9c Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 22:42:15 -0700
Subject: [PATCH 14/31] [Feature] UI - Teams: Modernize teams table with AntD,
 server-side pagination, and v2 API

Migrate the OldTeams table from Tremor to Ant Design components, matching the
Access Groups page pattern. Switch from /team/list to /v2/team/list for
server-side pagination, filtering, and sorting.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 ui/litellm-dashboard/src/app/page.tsx         |    6 +-
 .../src/components/OldTeams.test.tsx          |   66 +-
 .../src/components/OldTeams.tsx               | 1027 +++++++++--------
 .../TableIconActionButton.tsx                 |    2 +
 .../OrganizationDropdown.tsx                  |   12 +-
 .../src/components/ui/AntDLoadingSpinner.tsx  |   12 +
 6 files changed, 589 insertions(+), 536 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 5f2921203f..7147a5ea01 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -9,7 +9,7 @@ import AgentsPanel from "@/components/agents";
 import BudgetPanel from "@/components/budgets/budget_panel";
 import CacheDashboard from "@/components/cache_dashboard";
 import ClaudeCodePluginsPanel from "@/components/claude_code_plugins";
-import { fetchTeams } from "@/components/common_components/fetch_teams";
+import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams";
 import LoadingScreen from "@/components/common_components/LoadingScreen";
 import { CostTrackingSettings } from "@/components/CostTrackingSettings";
 import GeneralSettings from "@/components/general_settings";
@@ -339,7 +339,9 @@ function CreateKeyPageContent() {
       fetchUserModels(userID, userRole, accessToken, setUserModels);
     }
     if (accessToken && userID && userRole) {
-      fetchTeams(accessToken, userID, userRole, null, setTeams);
+      v2TeamListCall(accessToken, 1, 100, {
+        userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null,
+      }).then((response) => setTeams(response.teams ?? [])).catch(console.error);
     }
     if (accessToken) {
       fetchOrganizations(accessToken, setOrganizations);
diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx
index 7d64ac9afe..651d1495c6 100644
--- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx
+++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx
@@ -18,8 +18,8 @@ vi.mock("./networking", () => ({
   getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }),
 }));
 
-vi.mock("./common_components/fetch_teams", () => ({
-  fetchTeams: vi.fn(),
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  teamListCall: vi.fn().mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 0 }),
 }));
 
 vi.mock("./molecules/notifications_manager", () => ({
@@ -375,6 +375,9 @@ describe("OldTeams - handleCreate organization handling", () => {
         organizations={[]}
       />,
     );
+    await waitFor(() => {
+      expect(screen.getByTestId("delete-team-button")).toBeInTheDocument();
+    });
     const deleteTeamButton = screen.getByTestId("delete-team-button");
     act(() => {
       fireEvent.click(deleteTeamButton);
@@ -389,7 +392,7 @@ describe("OldTeams - empty state", () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
   });
 
-  it("should display empty state message when teams array is empty", () => {
+  it("should display empty state message when teams array is empty", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={[]}
@@ -402,11 +405,13 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.getByText("No teams found")).toBeInTheDocument();
-    expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("No teams yet")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Create your first team to organize members and manage access to models.")).toBeInTheDocument();
   });
 
-  it("should display empty state message when teams is null", () => {
+  it("should display empty state message when teams is null", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={null}
@@ -419,11 +424,13 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.getByText("No teams found")).toBeInTheDocument();
-    expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("No teams yet")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Create your first team to organize members and manage access to models.")).toBeInTheDocument();
   });
 
-  it("should not display empty state when teams array has items", () => {
+  it("should not display empty state when teams array has items", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={[
@@ -451,9 +458,11 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.queryByText("No teams found")).not.toBeInTheDocument();
-    expect(screen.queryByText("Adjust your filters or create a new team")).not.toBeInTheDocument();
-    expect(screen.getByText("Test Team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("Test Team")).toBeInTheDocument();
+    });
+    expect(screen.queryByText("No teams yet")).not.toBeInTheDocument();
+    expect(screen.queryByText("Create your first team to organize members and manage access to models.")).not.toBeInTheDocument();
   });
 });
 
@@ -621,12 +630,9 @@ describe("OldTeams - premium props", () => {
       />,
     );
 
-    const truncatedTeamId = "team-123456789".slice(0, 7);
-    const teamButton = await screen.findByRole("button", {
-      name: new RegExp(`${truncatedTeamId}\\.\\.\\.`),
-    });
+    const teamIdElement = await screen.findByText("team-123456789");
     act(() => {
-      fireEvent.click(teamButton);
+      fireEvent.click(teamIdElement);
     });
 
     await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled());
@@ -798,7 +804,7 @@ describe("OldTeams - access_group_ids in team create", () => {
       />,
     );
 
-    const createButton = screen.getByRole("button", { name: /create new team/i });
+    const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
     act(() => {
       fireEvent.click(createButton);
     });
@@ -823,7 +829,8 @@ describe("OldTeams - access_group_ids in team create", () => {
     const accessGroupInput = screen.getByTestId("access-group-selector");
     fireEvent.change(accessGroupInput, { target: { value: "ag-1,ag-2" } });
 
-    const createTeamSubmitButton = screen.getByRole("button", { name: /create team/i });
+    const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i });
+    const createTeamSubmitButton = createTeamSubmitButtons[createTeamSubmitButtons.length - 1];
     fireEvent.click(createTeamSubmitButton);
 
     await waitFor(() => {
@@ -865,7 +872,7 @@ describe("OldTeams - models dropdown options", () => {
       expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled();
     });
 
-    const createButton = screen.getByRole("button", { name: /create new team/i });
+    const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
     act(() => {
       fireEvent.click(createButton);
     });
@@ -884,7 +891,7 @@ describe("OldTeams - organization alias display", () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
   });
 
-  it("should display organization alias instead of organization id", () => {
+  it("should display organization alias instead of organization id", async () => {
     const mockOrganizations = [
       {
         organization_id: "org-123",
@@ -934,11 +941,13 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("Test Organization")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("Test Organization")).toBeInTheDocument();
+    });
     expect(screen.queryByText("org-123")).not.toBeInTheDocument();
   });
 
-  it("should display organization id when alias is not found", () => {
+  it("should display organization id when alias is not found", async () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
 
     renderWithQueryClient(
@@ -968,10 +977,12 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("org-unknown")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("org-unknown")).toBeInTheDocument();
+    });
   });
 
-  it("should display N/A when organization_id is null", () => {
+  it("should display N/A when organization_id is null", async () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
 
     renderWithQueryClient(
@@ -1001,6 +1012,9 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("N/A")).toBeInTheDocument();
+    await waitFor(() => {
+      // When organization_id is null, the table shows "—" in the Organization column
+      expect(screen.getAllByText("—").length).toBeGreaterThan(0);
+    });
   });
 });
diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx
index 702724c972..8970226b9f 100644
--- a/ui/litellm-dashboard/src/components/OldTeams.tsx
+++ b/ui/litellm-dashboard/src/components/OldTeams.tsx
@@ -3,50 +3,54 @@ import AvailableTeamsPanel from "@/components/team/available_teams";
 import TeamInfoView from "@/components/team/TeamInfo";
 import TeamSSOSettings from "@/components/TeamSSOSettings";
 import { isProxyAdminRole } from "@/utils/roles";
-import { InfoCircleOutlined } from "@ant-design/icons";
-import { ChevronDownIcon, ChevronRightIcon, RefreshIcon } from "@heroicons/react/outline";
-import { FilterInput } from "@/components/common_components/Filters/FilterInput";
-import { FiltersButton } from "@/components/common_components/Filters/FiltersButton";
-import { ResetFiltersButton } from "@/components/common_components/Filters/ResetFiltersButton";
-import { Search, User } from "lucide-react";
+import {
+  InfoCircleOutlined,
+  PlusOutlined,
+  TeamOutlined,
+  ReloadOutlined,
+} from "@ant-design/icons";
 import {
   Accordion,
   AccordionBody,
   AccordionHeader,
-  Badge,
-  Button,
-  Card,
-  Col,
-  Grid,
-  Icon,
-  Select,
-  SelectItem,
-  Tab,
-  TabGroup,
-  TabList,
-  TabPanel,
-  TabPanels,
-  Table,
-  TableBody,
-  TableCell,
-  TableHead,
-  TableHeaderCell,
-  TableRow,
-  Text,
   TextInput,
 } from "@tremor/react";
-import { Button as Button2, Form, Input, Modal, Select as Select2, Switch, Tooltip, Typography } from "antd";
-import React, { useEffect, useState } from "react";
-import { formatNumberWithCommas } from "../utils/dataUtils";
+import {
+  Button,
+  Card,
+  Flex,
+  Form,
+  Input,
+  Layout,
+  Modal,
+  Pagination,
+  Progress,
+  Select,
+  Space,
+  Switch,
+  Table,
+  Tabs,
+  Tag,
+  theme,
+  Tooltip,
+  Typography,
+  message,
+} from "antd";
+import type { ColumnsType } from "antd/es/table";
+import type { SorterResult } from "antd/es/table/interface";
+import { KeyIcon, LayersIcon, SearchIcon, UsersIcon } from "lucide-react";
+import React, { useEffect, useMemo, useRef, useState } from "react";
+import { AntDLoadingSpinner } from "@/components/ui/AntDLoadingSpinner";
+import OrganizationDropdown from "./common_components/OrganizationDropdown";
+import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
+import { teamListCall as v2TeamListCall, type TeamsResponse } from "@/app/(dashboard)/hooks/teams/useTeams";
 import AccessGroupSelector from "./common_components/AccessGroupSelector";
 import AgentSelector from "./agent_management/AgentSelector";
-import { fetchTeams } from "./common_components/fetch_teams";
 import ModelAliasManager from "./common_components/ModelAliasManager";
 import PremiumLoggingSettings from "./common_components/PremiumLoggingSettings";
 import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./common_components/RouterSettingsAccordion";
 import {
   fetchAvailableModelsForTeamOrKey,
-  getModelDisplayName,
   unfurlWildcardModelsInList,
 } from "./key_team_helpers/fetch_available_models_team_key";
 import type { KeyResponse, Team } from "./key_team_helpers/key_list";
@@ -85,8 +89,7 @@ interface EditTeamModalProps {
 
 import { updateExistingKeys } from "@/utils/dataUtils";
 import DeleteResourceModal from "./common_components/DeleteResourceModal";
-import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
-import { Member, teamCreateCall, v2TeamListCall } from "./networking";
+import { Member, teamCreateCall } from "./networking";
 import { ModelSelect } from "./ModelSelect/ModelSelect";
 
 interface TeamInfo {
@@ -182,10 +185,13 @@ const Teams: React.FC<TeamProps> = ({
 }) => {
   console.log(`organizations: ${JSON.stringify(organizations)}`);
   const { data: organizationsData } = useOrganizations();
-  const [lastRefreshed, setLastRefreshed] = useState("");
+  const [isLoading, setIsLoading] = useState(true);
+  const [fetchError, setFetchError] = useState<string | null>(null);
+  const [currentPage, setCurrentPage] = useState(1);
+  const [pageSize, setPageSize] = useState(10);
+  const [totalTeams, setTotalTeams] = useState(0);
   const [currentOrg, setCurrentOrg] = useState<Organization | null>(null);
   const [currentOrgForCreateTeam, setCurrentOrgForCreateTeam] = useState<Organization | null>(null);
-  const [showFilters, setShowFilters] = useState(false);
   const [filters, setFilters] = useState<FilterState>({
     team_id: "",
     team_alias: "",
@@ -193,19 +199,55 @@ const Teams: React.FC<TeamProps> = ({
     sort_by: "created_at",
     sort_order: "desc",
   });
+  const searchDebounceRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const [isSearching, setIsSearching] = useState(false);
+
+  const fetchTeamsV2 = async (opts: {
+    page?: number;
+    size?: number;
+    sortBy?: string;
+    sortOrder?: string;
+    organizationID?: string;
+    teamAlias?: string;
+  } = {}) => {
+    if (!accessToken) return;
+    const page = opts.page ?? currentPage;
+    const size = opts.size ?? pageSize;
+    const sortBy = opts.sortBy ?? filters.sort_by;
+    const sortOrder = opts.sortOrder ?? filters.sort_order;
+    const organizationID = opts.organizationID ?? filters.organization_id;
+    const teamAlias = opts.teamAlias ?? filters.team_alias;
+
+    setIsLoading(true);
+    setFetchError(null);
+    try {
+      const response: TeamsResponse = await v2TeamListCall(
+        accessToken,
+        page,
+        size,
+        {
+          organizationID: organizationID || null,
+          team_alias: teamAlias || null,
+          userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null,
+          sortBy: sortBy || null,
+          sortOrder: sortOrder || null,
+        },
+      );
+      setTeams(response.teams ?? []);
+      setTotalTeams(response.total ?? 0);
+    } catch (err: any) {
+      setFetchError(err?.message || "Failed to fetch teams");
+    } finally {
+      setIsLoading(false);
+    }
+  };
 
   useEffect(() => {
-    console.log(`inside useeffect - ${lastRefreshed}`);
-    if (accessToken) {
-      // Call your function here
-      fetchTeams(accessToken, userID, userRole, currentOrg, setTeams);
-    }
-    handleRefreshClick();
-  }, [lastRefreshed]);
+    fetchTeamsV2();
+  }, [accessToken]);
 
   const [form] = Form.useForm();
   const [memberForm] = Form.useForm();
-  const { Title, Paragraph } = Typography;
   const [value, setValue] = useState("");
   const [editModalVisible, setEditModalVisible] = useState(false);
 
@@ -225,7 +267,6 @@ const Teams: React.FC<TeamProps> = ({
   // Add this state near the other useState declarations
   const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
   const [policiesList, setPoliciesList] = useState<string[]>([]);
-  const [expandedAccordions, setExpandedAccordions] = useState<Record<string, boolean>>({});
   const [loggingSettings, setLoggingSettings] = useState<any[]>([]);
   const [mcpAccessGroups, setMcpAccessGroups] = useState<string[]>([]);
   const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false);
@@ -376,7 +417,7 @@ const Teams: React.FC<TeamProps> = ({
     try {
       setIsTeamDeleting(true);
       await teamDeleteCall(accessToken, teamToDelete.team_id);
-      await fetchTeams(accessToken, userID, userRole, currentOrg, setTeams);
+      await fetchTeamsV2();
       NotificationsManager.success("Team deleted successfully");
     } catch (error) {
       NotificationsManager.fromBackend("Error deleting the team: " + error);
@@ -572,467 +613,447 @@ const Teams: React.FC<TeamProps> = ({
     return false;
   };
 
-  const handleRefreshClick = () => {
-    // Update the 'lastRefreshed' state to the current date and time
-    const currentDate = new Date();
-    setLastRefreshed(currentDate.toLocaleString());
+  const handleSearchChange = (value: string) => {
+    if (searchDebounceRef.current) clearTimeout(searchDebounceRef.current);
+    setIsSearching(true);
+    searchDebounceRef.current = setTimeout(async () => {
+      try {
+        setFilters((prev) => ({ ...prev, team_alias: value }));
+        setCurrentPage(1);
+        await fetchTeamsV2({ page: 1, teamAlias: value });
+      } finally {
+        setIsSearching(false);
+      }
+    }, 300);
   };
 
-  const handleFilterChange = (key: keyof FilterState, value: string) => {
+  const handleFilterChange = async (key: keyof FilterState, value: string) => {
     const newFilters = { ...filters, [key]: value };
     setFilters(newFilters);
-    // Call teamListCall with the new filters
-    if (accessToken) {
-      v2TeamListCall(
+    setCurrentPage(1);
+    if (!accessToken) return;
+    try {
+      const response: TeamsResponse = await v2TeamListCall(
         accessToken,
-        newFilters.organization_id || null,
-        null,
-        newFilters.team_id || null,
-        newFilters.team_alias || null,
-      )
-        .then((response) => {
-          if (response && response.teams) {
-            setTeams(response.teams);
-          }
-        })
-        .catch((error) => {
-          console.error("Error fetching teams:", error);
-        });
-    }
-  };
-
-  const handleSortChange = (sortBy: string, sortOrder: "asc" | "desc") => {
-    const newFilters = {
-      ...filters,
-      sort_by: sortBy,
-      sort_order: sortOrder,
-    };
-    setFilters(newFilters);
-    // Call teamListCall with the new sort parameters
-    if (accessToken) {
-      v2TeamListCall(
-        accessToken,
-        filters.organization_id || null,
-        null,
-        filters.team_id || null,
-        filters.team_alias || null,
-      )
-        .then((response) => {
-          if (response && response.teams) {
-            setTeams(response.teams);
-          }
-        })
-        .catch((error) => {
-          console.error("Error fetching teams:", error);
-        });
+        1,
+        pageSize,
+        {
+          organizationID: newFilters.organization_id || null,
+          team_alias: newFilters.team_alias || null,
+          userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null,
+          sortBy: newFilters.sort_by || null,
+          sortOrder: newFilters.sort_order || null,
+        },
+      );
+      setTeams(response.teams ?? []);
+      setTotalTeams(response.total ?? 0);
+    } catch (error) {
+      console.error("Error fetching teams:", error);
     }
   };
 
   const handleFilterReset = () => {
-    setFilters({
+    if (searchDebounceRef.current) clearTimeout(searchDebounceRef.current);
+    setIsSearching(false);
+    const resetFilters: FilterState = {
       team_id: "",
       team_alias: "",
       organization_id: "",
       sort_by: "created_at",
       sort_order: "desc",
-    });
-    // Reset teams list
-    if (accessToken) {
-      v2TeamListCall(accessToken, null, userID || null, null, null)
-        .then((response) => {
-          if (response && response.teams) {
-            setTeams(response.teams);
-          }
-        })
-        .catch((error) => {
-          console.error("Error fetching teams:", error);
-        });
-    }
+    };
+    setFilters(resetFilters);
+    setCurrentPage(1);
+    fetchTeamsV2({ page: 1, organizationID: "", teamAlias: "", sortBy: "created_at", sortOrder: "desc" });
   };
 
+  const { token } = theme.useToken();
+  const { Title, Text } = Typography;
+  const { Content } = Layout;
+
+  const handleRetry = () => {
+    fetchTeamsV2();
+  };
+
+  const handleTableSort = (_pagination: unknown, _filters: unknown, sorter: SorterResult<Team> | SorterResult<Team>[]) => {
+    const s = Array.isArray(sorter) ? sorter[0] : sorter;
+    const sortBy = s.order ? (s.columnKey as string) : "created_at";
+    const sortOrder = s.order === "ascend" ? "asc" : s.order === "descend" ? "desc" : "desc";
+    setFilters((prev) => ({ ...prev, sort_by: sortBy, sort_order: sortOrder }));
+    fetchTeamsV2({ sortBy, sortOrder });
+  };
+
+  const teamColumns: ColumnsType<Team> = useMemo(() => [
+    {
+      title: "Team ID",
+      dataIndex: "team_id",
+      key: "team_id",
+      width: 170,
+      ellipsis: true,
+      render: (id: string, record: Team) => (
+        <Tooltip title={id}>
+          <Text
+            ellipsis
+            className="text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs cursor-pointer"
+            style={{ fontSize: 14, padding: "1px 8px" }}
+            onClick={() => setSelectedTeamId(record.team_id)}
+          >
+            {id}
+          </Text>
+        </Tooltip>
+      ),
+    },
+    {
+      title: "Team Alias",
+      dataIndex: "team_alias",
+      key: "team_alias",
+      ellipsis: true,
+      sorter: true,
+      render: (alias: string | undefined) => (
+        <Text style={{ fontSize: 14 }}>
+          {alias || <Text type="secondary" italic>—</Text>}
+        </Text>
+      ),
+    },
+    {
+      title: "Organization",
+      key: "organization",
+      width: 160,
+      ellipsis: true,
+      render: (_: unknown, record: Team) => {
+        const orgAlias = getOrganizationAlias(record.organization_id, organizationsData || organizations);
+        return record.organization_id ? <Text ellipsis style={{ fontSize: 14 }}>{orgAlias}</Text> : <Text type="secondary">—</Text>;
+      },
+    },
+    {
+      title: "Resources",
+      key: "resources",
+      width: 240,
+      render: (_: unknown, record: Team) => {
+        const memberCount = perTeamInfo?.[record.team_id]?.team_info?.members_with_roles?.length ?? 0;
+        const modelCount = record.models?.length ?? 0;
+        const keyCount = perTeamInfo?.[record.team_id]?.keys?.length ?? 0;
+        return (
+          <Flex gap={12} align="center">
+            <Tooltip title={`${memberCount} Members`}>
+              <Tag color="purple" style={{ fontSize: 14, padding: "2px 8px", margin: 0 }}>
+                <Flex align="center" gap={6}>
+                  <UsersIcon size={14} />
+                  {memberCount}
+                </Flex>
+              </Tag>
+            </Tooltip>
+            <Tooltip title={`${modelCount} Models`}>
+              <Tag color="blue" style={{ fontSize: 14, padding: "2px 8px", margin: 0 }}>
+                <Flex align="center" gap={6}>
+                  <LayersIcon size={14} />
+                  {modelCount}
+                </Flex>
+              </Tag>
+            </Tooltip>
+            <Tooltip title={`${keyCount} Keys`}>
+              <Tag color="cyan" style={{ fontSize: 14, padding: "2px 8px", margin: 0 }}>
+                <Flex align="center" gap={6}>
+                  <KeyIcon size={14} />
+                  {keyCount}
+                </Flex>
+              </Tag>
+            </Tooltip>
+          </Flex>
+        );
+      },
+    },
+    {
+      title: "Spend / Budget",
+      key: "spend",
+      width: 200,
+      sorter: true,
+      render: (_: unknown, record: Team) => {
+        const spendVal = record.spend ?? 0;
+        const budgetVal = record.max_budget;
+        const spendStr = `$${spendVal.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`;
+        const budgetStr = budgetVal != null
+          ? `$${budgetVal.toLocaleString(undefined, { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
+          : "Unlimited";
+        const percent = budgetVal != null && budgetVal > 0 ? Math.min((spendVal / budgetVal) * 100, 100) : null;
+        return (
+          <Flex vertical gap={2}>
+            <Text style={{ fontSize: 13 }}>
+              {spendStr}
+              <Text type="secondary" style={{ fontSize: 12 }}>{" / "}{budgetStr}</Text>
+            </Text>
+            {percent != null && (
+              <Progress
+                percent={percent}
+                size="small"
+                showInfo={false}
+                strokeColor={percent >= 90 ? "#ff4d4f" : percent >= 70 ? "#faad14" : "#1677ff"}
+                style={{ marginBottom: 0 }}
+              />
+            )}
+          </Flex>
+        );
+      },
+    },
+    {
+      title: "Created",
+      dataIndex: "created_at",
+      key: "created_at",
+      width: 130,
+      ellipsis: true,
+      sorter: true,
+      render: (date: string | undefined) => (
+        <Text type="secondary" style={{ fontSize: 13 }}>
+          {date ? new Date(date).toLocaleDateString(undefined, { year: "numeric", month: "short", day: "numeric" }) : "—"}
+        </Text>
+      ),
+    },
+    {
+      title: "Actions",
+      key: "actions",
+      width: 120,
+      align: "right" as const,
+      render: (_: unknown, record: Team) => (
+        <Space size={4}>
+          <TableIconActionButton
+            variant="Copy"
+            tooltipText="Copy Team ID"
+            onClick={() => {
+              navigator.clipboard.writeText(record.team_id)
+                .then(() => message.success("Team ID copied"))
+                .catch(() => message.error("Failed to copy"));
+            }}
+          />
+          {userRole === "Admin" && (
+            <>
+              <TableIconActionButton
+                variant="Edit"
+                tooltipText="Edit team"
+                dataTestId="edit-team-button"
+                onClick={() => {
+                  setSelectedTeamId(record.team_id);
+                  setEditTeam(true);
+                }}
+              />
+              <TableIconActionButton
+                variant="Delete"
+                tooltipText="Delete team"
+                dataTestId="delete-team-button"
+                onClick={() => handleDelete(record)}
+              />
+            </>
+          )}
+        </Space>
+      ),
+    },
+  ], [userRole, perTeamInfo, organizationsData, organizations]);
+
+  const displayTeams = useMemo(() => teams ?? [], [teams]);
+
+  const renderTeamsContent = () => {
+    if (isLoading) {
+      return (
+        <Flex justify="center" align="center" style={{ padding: "80px 0" }}>
+          <AntDLoadingSpinner fontSize={48} />
+        </Flex>
+      );
+    }
+
+    if (fetchError) {
+      return (
+        <Flex vertical align="center" gap={16} style={{ padding: "64px 0" }}>
+          <Text type="danger" style={{ fontSize: 15 }}>
+            Failed to load teams
+          </Text>
+          <Text type="secondary" style={{ fontSize: 13 }}>
+            {fetchError}
+          </Text>
+          <Button icon={<ReloadOutlined />} onClick={handleRetry}>
+            Retry
+          </Button>
+        </Flex>
+      );
+    }
+
+    return (
+      <Table<Team>
+        columns={teamColumns}
+        dataSource={displayTeams}
+        rowKey="team_id"
+        pagination={false}
+        onChange={handleTableSort}
+        locale={{
+          emptyText: (
+            <div style={{ padding: "64px 0", textAlign: "center" }}>
+              <TeamOutlined style={{ fontSize: 40, color: "#d9d9d9", marginBottom: 12 }} />
+              <div>
+                <Text style={{ fontSize: 15, color: "#595959" }}>No teams yet</Text>
+              </div>
+              <div style={{ marginTop: 4 }}>
+                <Text type="secondary" style={{ fontSize: 13 }}>
+                  Create your first team to organize members and manage access to models.
+                </Text>
+              </div>
+              {canCreateOrManageTeams(userRole, userID, organizations) && (
+                <Button
+                  type="primary"
+                  icon={<PlusOutlined />}
+                  onClick={() => setIsTeamModalVisible(true)}
+                  style={{ marginTop: 16 }}
+                >
+                  Create Team
+                </Button>
+              )}
+            </div>
+          ),
+        }}
+        scroll={{ x: 1000 }}
+        size="middle"
+      />
+    );
+  };
+
+  const tabItems = [
+    {
+      key: "your-teams",
+      label: "Your Teams",
+      children: (
+        <>
+          <Card styles={{ body: { padding: 0 } }}>
+            <Flex
+              justify="space-between"
+              align="center"
+              style={{ padding: "12px 16px" }}
+            >
+              <Flex gap={12} align="center">
+                <Input
+                  prefix={<SearchIcon size={16} />}
+                  suffix={isSearching ? <AntDLoadingSpinner size="small" /> : null}
+                  placeholder="Search teams by name..."
+                  onChange={(e) => handleSearchChange(e.target.value)}
+                  allowClear
+                  style={{ maxWidth: 400 }}
+                />
+                <OrganizationDropdown
+                  organizations={organizations}
+                  value={filters.organization_id || undefined}
+                  onChange={(value: string) => handleFilterChange("organization_id", value || "")}
+                  loading={isLoading}
+                />
+              </Flex>
+              <Pagination
+                current={currentPage}
+                total={totalTeams}
+                pageSize={pageSize}
+                onChange={(page, size) => {
+                  setCurrentPage(page);
+                  setPageSize(size);
+                  fetchTeamsV2({ page, size });
+                }}
+                size="small"
+                showTotal={(total) => `${total} teams`}
+                showSizeChanger
+                pageSizeOptions={["10", "20", "50"]}
+              />
+            </Flex>
+
+            {renderTeamsContent()}
+          </Card>
+
+          <DeleteResourceModal
+            isOpen={isDeleteModalOpen}
+            title="Delete Team?"
+            alertMessage={
+              teamToDelete?.keys?.length === 0
+                ? undefined
+                : `Warning: This team has ${teamToDelete?.keys?.length} keys associated with it. Deleting the team will also delete all associated keys. This action is irreversible.`
+            }
+            message="Are you sure you want to delete this team and all its keys? This action cannot be undone."
+            resourceInformationTitle="Team Information"
+            resourceInformation={[
+              { label: "Team ID", value: teamToDelete?.team_id, code: true },
+              { label: "Team Name", value: teamToDelete?.team_alias },
+              { label: "Keys", value: teamToDelete?.keys?.length },
+              { label: "Members", value: teamToDelete?.members_with_roles?.length },
+            ]}
+            requiredConfirmation={teamToDelete?.team_alias}
+            onCancel={cancelDelete}
+            onOk={confirmDelete}
+            confirmLoading={isTeamDeleting}
+          />
+        </>
+      ),
+    },
+    {
+      key: "available-teams",
+      label: "Available Teams",
+      children: <AvailableTeamsPanel accessToken={accessToken} userID={userID} />,
+    },
+    ...(isProxyAdminRole(userRole || "")
+      ? [
+          {
+            key: "default-settings",
+            label: "Default Team Settings",
+            children: <TeamSSOSettings accessToken={accessToken} userID={userID || ""} userRole={userRole || ""} />,
+          },
+        ]
+      : []),
+  ];
+
   return (
-    <div className="w-full mx-4 h-[75vh]">
-      <Grid numItems={1} className="gap-2 p-8 w-full mt-2">
-        <Col numColSpan={1} className="flex flex-col gap-2">
-          {canCreateOrManageTeams(userRole, userID, organizations) && (
-            <Button className="w-fit" onClick={() => setIsTeamModalVisible(true)}>
-              + Create New Team
-            </Button>
-          )}
-          {selectedTeamId ? (
-            <TeamInfoView
-              teamId={selectedTeamId}
-              onUpdate={(data) => {
-                setTeams((teams) => {
-                  if (teams == null) {
-                    return teams;
-                  }
-                  const updated = teams.map((team) => {
-                    if (data.team_id === team.team_id) {
-                      return updateExistingKeys(team, data);
-                    }
-                    return team;
-                  });
-                  // Minimal fix: refresh the full team list after an update
-                  if (accessToken) {
-                    fetchTeams(accessToken, userID, userRole, currentOrg, setTeams);
-                  }
-                  return updated;
-                });
-              }}
-              onClose={() => {
-                setSelectedTeamId(null);
-                setEditTeam(false);
-              }}
-              accessToken={accessToken}
-              is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))}
-              is_proxy_admin={userRole == "Admin"}
-              userModels={userModels}
-              editTeam={editTeam}
-              premiumUser={premiumUser}
-            />
-          ) : (
-            <TabGroup className="gap-2 h-[75vh] w-full">
-              <TabList className="flex justify-between mt-2 w-full items-center">
-                <div className="flex">
-                  <Tab>Your Teams</Tab>
-                  <Tab>Available Teams</Tab>
-                  {isProxyAdminRole(userRole || "") && <Tab>Default Team Settings</Tab>}
-                </div>
-                <div className="flex items-center space-x-2">
-                  {lastRefreshed && <Text>Last Refreshed: {lastRefreshed}</Text>}
-                  <Icon
-                    icon={RefreshIcon} // Modify as necessary for correct icon name
-                    variant="shadow"
-                    size="xs"
-                    className="self-center"
-                    onClick={handleRefreshClick}
-                  />
-                </div>
-              </TabList>
-              <TabPanels>
-                <TabPanel>
-                  <Text>
-                    Click on &ldquo;Team ID&rdquo; to view team details <b>and</b> manage team members.
-                  </Text>
-                  <Grid numItems={1} className="gap-2 pt-2 pb-2 h-[75vh] w-full mt-2">
-                    <Col numColSpan={1}>
-                      <Card className="w-full mx-auto flex-auto overflow-hidden overflow-y-auto max-h-[50vh]">
-                        <div className="border-b px-6 py-4">
-                          <div className="flex flex-col space-y-4">
-                            {/* Search and Filter Controls */}
-                            <div className="flex flex-wrap items-center gap-3">
-                              {/* Team Alias Search */}
-                              <FilterInput
-                                placeholder="Search by Team Name..."
-                                value={filters.team_alias}
-                                onChange={(value) => handleFilterChange("team_alias", value)}
-                                icon={Search}
-                              />
+    <Content style={{ padding: token.paddingLG, paddingInline: token.paddingLG * 2 }}>
+      {selectedTeamId ? (
+        <TeamInfoView
+          teamId={selectedTeamId}
+          onUpdate={(data) => {
+            setTeams((teams) => {
+              if (teams == null) {
+                return teams;
+              }
+              return teams.map((team) => {
+                if (data.team_id === team.team_id) {
+                  return updateExistingKeys(team, data);
+                }
+                return team;
+              });
+            });
+            fetchTeamsV2();
+          }}
+          onClose={() => {
+            setSelectedTeamId(null);
+            setEditTeam(false);
+          }}
+          accessToken={accessToken}
+          is_team_admin={is_team_admin(teams?.find((team) => team.team_id === selectedTeamId))}
+          is_proxy_admin={userRole == "Admin"}
+          userModels={userModels}
+          editTeam={editTeam}
+          premiumUser={premiumUser}
+        />
+      ) : (
+        <>
+          <Flex justify="space-between" align="center" style={{ marginBottom: 16 }}>
+            <Space direction="vertical" size={0}>
+              <Title level={2} style={{ margin: 0 }}>
+                <TeamOutlined style={{ marginRight: 8 }} />
+                Teams
+              </Title>
+              <Text type="secondary">
+                Manage teams, members, and their access to models and budgets
+              </Text>
+            </Space>
+            {canCreateOrManageTeams(userRole, userID, organizations) && (
+              <Button type="primary" icon={<PlusOutlined />} onClick={() => setIsTeamModalVisible(true)}>
+                Create Team
+              </Button>
+            )}
+          </Flex>
 
-                              {/* Filter Button */}
-                              <FiltersButton
-                                onClick={() => setShowFilters(!showFilters)}
-                                active={showFilters}
-                                hasActiveFilters={!!(filters.team_id || filters.team_alias || filters.organization_id)}
-                              />
+          <Tabs items={tabItems} />
+        </>
+      )}
 
-                              {/* Reset Filters Button */}
-                              <ResetFiltersButton onClick={handleFilterReset} />
-                            </div>
-
-                            {/* Additional Filters */}
-                            {showFilters && (
-                              <div className="flex flex-wrap items-center gap-3 mt-3">
-                                {/* Team ID Search */}
-                                <FilterInput
-                                  placeholder="Enter Team ID"
-                                  value={filters.team_id}
-                                  onChange={(value) => handleFilterChange("team_id", value)}
-                                  icon={User}
-                                />
-
-                                {/* Organization Dropdown */}
-                                <div className="w-64">
-                                  <Select
-                                    value={filters.organization_id || ""}
-                                    onValueChange={(value) => handleFilterChange("organization_id", value)}
-                                    placeholder="Select Organization"
-                                  >
-                                    {organizations?.map((org) => (
-                                      <SelectItem key={org.organization_id} value={org.organization_id || ""}>
-                                        {org.organization_alias || org.organization_id}
-                                      </SelectItem>
-                                    ))}
-                                  </Select>
-                                </div>
-                              </div>
-                            )}
-                          </div>
-                        </div>
-                        <Table>
-                          <TableHead>
-                            <TableRow>
-                              <TableHeaderCell>Team Name</TableHeaderCell>
-                              <TableHeaderCell>Team ID</TableHeaderCell>
-                              <TableHeaderCell>Created</TableHeaderCell>
-                              <TableHeaderCell>Spend (USD)</TableHeaderCell>
-                              <TableHeaderCell>Budget (USD)</TableHeaderCell>
-                              <TableHeaderCell>Models</TableHeaderCell>
-                              <TableHeaderCell>Organization</TableHeaderCell>
-                              <TableHeaderCell>Info</TableHeaderCell>
-                              <TableHeaderCell>Actions</TableHeaderCell>
-                            </TableRow>
-                          </TableHead>
-
-                          <TableBody>
-                            {teams && teams.length > 0 ? (
-                              teams
-                                .filter((team) => {
-                                  if (!currentOrg) return true;
-                                  return team.organization_id === currentOrg.organization_id;
-                                })
-                                .sort((a, b) => new Date(b.created_at).getTime() - new Date(a.created_at).getTime())
-                                .map((team: any) => (
-                                  <TableRow key={team.team_id}>
-                                    <TableCell
-                                      style={{
-                                        maxWidth: "4px",
-                                        whiteSpace: "pre-wrap",
-                                        overflow: "hidden",
-                                      }}
-                                    >
-                                      {team["team_alias"]}
-                                    </TableCell>
-                                    <TableCell>
-                                      <div className="overflow-hidden">
-                                        <Tooltip title={team.team_id}>
-                                          <Button
-                                            size="xs"
-                                            variant="light"
-                                            className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]"
-                                            onClick={() => {
-                                              // Add click handler
-                                              setSelectedTeamId(team.team_id);
-                                            }}
-                                          >
-                                            {team.team_id.slice(0, 7)}...
-                                          </Button>
-                                        </Tooltip>
-                                      </div>
-                                    </TableCell>
-                                    <TableCell
-                                      style={{
-                                        maxWidth: "4px",
-                                        whiteSpace: "pre-wrap",
-                                        overflow: "hidden",
-                                      }}
-                                    >
-                                      {team.created_at ? new Date(team.created_at).toLocaleDateString() : "N/A"}
-                                    </TableCell>
-                                    <TableCell
-                                      style={{
-                                        maxWidth: "4px",
-                                        whiteSpace: "pre-wrap",
-                                        overflow: "hidden",
-                                      }}
-                                    >
-                                      {formatNumberWithCommas(team["spend"], 4)}
-                                    </TableCell>
-                                    <TableCell
-                                      style={{
-                                        maxWidth: "4px",
-                                        whiteSpace: "pre-wrap",
-                                        overflow: "hidden",
-                                      }}
-                                    >
-                                      {team["max_budget"] !== null && team["max_budget"] !== undefined
-                                        ? team["max_budget"]
-                                        : "No limit"}
-                                    </TableCell>
-                                    <TableCell
-                                      style={{
-                                        maxWidth: "8-x",
-                                        whiteSpace: "pre-wrap",
-                                        overflow: "hidden",
-                                      }}
-                                      className={team.models.length > 3 ? "px-0" : ""}
-                                    >
-                                      <div className="flex flex-col">
-                                        {Array.isArray(team.models) ? (
-                                          <div className="flex flex-col">
-                                            {team.models.length === 0 ? (
-                                              <Badge size={"xs"} className="mb-1" color="red">
-                                                <Text>All Proxy Models</Text>
-                                              </Badge>
-                                            ) : (
-                                              <>
-                                                <div className="flex items-start">
-                                                  {team.models.length > 3 && (
-                                                    <div>
-                                                      <Icon
-                                                        icon={
-                                                          expandedAccordions[team.team_id]
-                                                            ? ChevronDownIcon
-                                                            : ChevronRightIcon
-                                                        }
-                                                        className="cursor-pointer"
-                                                        size="xs"
-                                                        onClick={() => {
-                                                          setExpandedAccordions((prev) => ({
-                                                            ...prev,
-                                                            [team.team_id]: !prev[team.team_id],
-                                                          }));
-                                                        }}
-                                                      />
-                                                    </div>
-                                                  )}
-                                                  <div className="flex flex-wrap gap-1">
-                                                    {team.models.slice(0, 3).map((model: string, index: number) =>
-                                                      model === "all-proxy-models" ? (
-                                                        <Badge key={index} size={"xs"} color="red">
-                                                          <Text>All Proxy Models</Text>
-                                                        </Badge>
-                                                      ) : (
-                                                        <Badge key={index} size={"xs"} color="blue">
-                                                          <Text>
-                                                            {model.length > 30
-                                                              ? `${getModelDisplayName(model).slice(0, 30)}...`
-                                                              : getModelDisplayName(model)}
-                                                          </Text>
-                                                        </Badge>
-                                                      ),
-                                                    )}
-                                                    {team.models.length > 3 && !expandedAccordions[team.team_id] && (
-                                                      <Badge size={"xs"} color="gray" className="cursor-pointer">
-                                                        <Text>
-                                                          +{team.models.length - 3}{" "}
-                                                          {team.models.length - 3 === 1 ? "more model" : "more models"}
-                                                        </Text>
-                                                      </Badge>
-                                                    )}
-                                                    {expandedAccordions[team.team_id] && (
-                                                      <div className="flex flex-wrap gap-1">
-                                                        {team.models.slice(3).map((model: string, index: number) =>
-                                                          model === "all-proxy-models" ? (
-                                                            <Badge key={index + 3} size={"xs"} color="red">
-                                                              <Text>All Proxy Models</Text>
-                                                            </Badge>
-                                                          ) : (
-                                                            <Badge key={index + 3} size={"xs"} color="blue">
-                                                              <Text>
-                                                                {model.length > 30
-                                                                  ? `${getModelDisplayName(model).slice(0, 30)}...`
-                                                                  : getModelDisplayName(model)}
-                                                              </Text>
-                                                            </Badge>
-                                                          ),
-                                                        )}
-                                                      </div>
-                                                    )}
-                                                  </div>
-                                                </div>
-                                              </>
-                                            )}
-                                          </div>
-                                        ) : null}
-                                      </div>
-                                    </TableCell>
-
-                                    <TableCell>
-                                      {getOrganizationAlias(team.organization_id, organizationsData || organizations)}
-                                    </TableCell>
-                                    <TableCell>
-                                      <Text>
-                                        {perTeamInfo &&
-                                          team.team_id &&
-                                          perTeamInfo[team.team_id] &&
-                                          perTeamInfo[team.team_id].keys &&
-                                          perTeamInfo[team.team_id].keys.length}{" "}
-                                        Keys
-                                      </Text>
-                                      <Text>
-                                        {perTeamInfo &&
-                                          team.team_id &&
-                                          perTeamInfo[team.team_id] &&
-                                          perTeamInfo[team.team_id].team_info &&
-                                          perTeamInfo[team.team_id].team_info.members_with_roles &&
-                                          perTeamInfo[team.team_id].team_info.members_with_roles.length}{" "}
-                                        Members
-                                      </Text>
-                                    </TableCell>
-                                    <TableCell>
-                                      {userRole == "Admin" ? (
-                                        <>
-                                          <TableIconActionButton
-                                            variant="Edit"
-                                            onClick={() => {
-                                              setSelectedTeamId(team.team_id);
-                                              setEditTeam(true);
-                                            }}
-                                            dataTestId="edit-team-button"
-                                            tooltipText="Edit team"
-                                          />
-                                          <TableIconActionButton
-                                            variant="Delete"
-                                            onClick={() => handleDelete(team)}
-                                            dataTestId="delete-team-button"
-                                            tooltipText="Delete team"
-                                          />
-                                        </>
-                                      ) : null}
-                                    </TableCell>
-                                  </TableRow>
-                                ))
-                            ) : (
-                              <TableRow>
-                                <TableCell colSpan={9} className="text-center">
-                                  <div className="flex flex-col items-center justify-center py-4">
-                                    <Text className="text-lg font-medium mb-2">No teams found</Text>
-                                    <Text className="text-sm">Adjust your filters or create a new team</Text>
-                                  </div>
-                                </TableCell>
-                              </TableRow>
-                            )}
-                          </TableBody>
-                        </Table>
-                        <DeleteResourceModal
-                          isOpen={isDeleteModalOpen}
-                          title="Delete Team?"
-                          alertMessage={
-                            teamToDelete?.keys?.length === 0
-                              ? undefined
-                              : `Warning: This team has ${teamToDelete?.keys?.length} keys associated with it. Deleting the team will also delete all associated keys. This action is irreversible.`
-                          }
-                          message="Are you sure you want to delete this team and all its keys? This action cannot be undone."
-                          resourceInformationTitle="Team Information"
-                          resourceInformation={[
-                            { label: "Team ID", value: teamToDelete?.team_id, code: true },
-                            { label: "Team Name", value: teamToDelete?.team_alias },
-                            { label: "Keys", value: teamToDelete?.keys?.length },
-                            { label: "Members", value: teamToDelete?.members_with_roles?.length },
-                          ]}
-                          requiredConfirmation={teamToDelete?.team_alias}
-                          onCancel={cancelDelete}
-                          onOk={confirmDelete}
-                          confirmLoading={isTeamDeleting}
-                        />
-                      </Card>
-                    </Col>
-                  </Grid>
-                </TabPanel>
-                <TabPanel>
-                  <AvailableTeamsPanel accessToken={accessToken} userID={userID} />
-                </TabPanel>
-                {isProxyAdminRole(userRole || "") && (
-                  <TabPanel>
-                    <TeamSSOSettings accessToken={accessToken} userID={userID || ""} userRole={userRole || ""} />
-                  </TabPanel>
-                )}
-              </TabPanels>
-            </TabGroup>
-          )}
-          {canCreateOrManageTeams(userRole, userID, organizations) && (
+      {canCreateOrManageTeams(userRole, userID, organizations) && (
             <Modal
               title="Create Team"
               open={isTeamModalVisible}
@@ -1117,7 +1138,7 @@ const Teams: React.FC<TeamProps> = ({
                                 : ""
                           }
                         >
-                          <Select2
+                          <Select
                             showSearch
                             allowClear={!isOrgAdmin}
                             disabled={isSingleOrg}
@@ -1136,18 +1157,18 @@ const Teams: React.FC<TeamProps> = ({
                             optionFilterProp="children"
                           >
                             {adminOrgs?.map((org) => (
-                              <Select2.Option key={org.organization_id} value={org.organization_id}>
+                              <Select.Option key={org.organization_id} value={org.organization_id}>
                                 <span className="font-medium">{org.organization_alias}</span>{" "}
                                 <span className="text-gray-500">({org.organization_id})</span>
-                              </Select2.Option>
+                              </Select.Option>
                             ))}
-                          </Select2>
+                          </Select>
                         </Form.Item>
 
                         {/* Show message when org admin needs to select organization */}
                         {isOrgAdmin && !isSingleOrg && adminOrgs.length > 1 && (
                           <div className="mb-8 p-4 bg-blue-50 border border-blue-200 rounded-md">
-                            <Text className="text-blue-800 text-sm">
+                            <Text style={{ color: "#1e40af", fontSize: 14 }}>
                               Please select an organization to create a team for. You can only create teams within
                               organizations where you are an admin.
                             </Text>
@@ -1190,11 +1211,11 @@ const Teams: React.FC<TeamProps> = ({
                     <NumericalInput step={0.01} precision={2} width={200} />
                   </Form.Item>
                   <Form.Item className="mt-8" label="Reset Budget" name="budget_duration">
-                    <Select2 defaultValue={null} placeholder="n/a">
-                      <Select2.Option value="24h">daily</Select2.Option>
-                      <Select2.Option value="7d">weekly</Select2.Option>
-                      <Select2.Option value="30d">monthly</Select2.Option>
-                    </Select2>
+                    <Select defaultValue={null} placeholder="n/a">
+                      <Select.Option value="24h">daily</Select.Option>
+                      <Select.Option value="7d">weekly</Select.Option>
+                      <Select.Option value="30d">monthly</Select.Option>
+                    </Select>
                   </Form.Item>
                   <Form.Item label="Tokens per minute Limit (TPM)" name="tpm_limit">
                     <NumericalInput step={1} width={400} />
@@ -1313,7 +1334,7 @@ const Teams: React.FC<TeamProps> = ({
                         className="mt-8"
                         help="Select existing guardrails or enter new ones"
                       >
-                        <Select2
+                        <Select
                           mode="tags"
                           style={{ width: "100%" }}
                           placeholder="Select or enter guardrails"
@@ -1367,7 +1388,7 @@ const Teams: React.FC<TeamProps> = ({
                         className="mt-8"
                         help="Select existing policies or enter new ones"
                       >
-                        <Select2
+                        <Select
                           mode="tags"
                           style={{ width: "100%" }}
                           placeholder="Select or enter policies"
@@ -1533,7 +1554,7 @@ const Teams: React.FC<TeamProps> = ({
                     </AccordionHeader>
                     <AccordionBody>
                       <div className="mt-4">
-                        <Text className="text-sm text-gray-600 mb-4">
+                        <Text type="secondary" style={{ fontSize: 14, marginBottom: 16, display: "block" }}>
                           Create custom aliases for models that can be used by team members in API calls. This allows
                           you to create shortcuts for specific models.
                         </Text>
@@ -1548,14 +1569,12 @@ const Teams: React.FC<TeamProps> = ({
                   </Accordion>
                 </>
                 <div style={{ textAlign: "right", marginTop: "10px" }}>
-                  <Button2 htmlType="submit">Create Team</Button2>
+                  <Button htmlType="submit">Create Team</Button>
                 </div>
               </Form>
             </Modal>
           )}
-        </Col>
-      </Grid>
-    </div>
+    </Content>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx
index 488913a734..2f146aab72 100644
--- a/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx
@@ -6,6 +6,7 @@ import {
   ChevronUpIcon,
   ChevronDownIcon,
   ExternalLinkIcon,
+  ClipboardCopyIcon,
 } from "@heroicons/react/outline";
 import { Tooltip } from "antd";
 import BaseActionButton from "../BaseActionButton";
@@ -32,6 +33,7 @@ export const TableIconActionButtonMap: Record<string, TableIconActionButtonBaseP
   Up: { icon: ChevronUpIcon, className: "hover:text-blue-600" },
   Down: { icon: ChevronDownIcon, className: "hover:text-blue-600" },
   Open: { icon: ExternalLinkIcon, className: "hover:text-green-600" },
+  Copy: { icon: ClipboardCopyIcon, className: "hover:text-blue-600" },
 };
 
 export default function TableIconActionButton({
diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx
index ac93041f3d..db8e75302f 100644
--- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx
@@ -1,13 +1,16 @@
 import React from "react";
-import { Select } from "antd";
+import { Select, Typography } from "antd";
 import { Organization } from "../networking";
 
+const { Text } = Typography;
+
 interface OrganizationDropdownProps {
   organizations?: Organization[] | null;
   value?: string;
   onChange?: (value: string) => void;
   disabled?: boolean;
   loading?: boolean;
+  style?: React.CSSProperties;
 }
 
 const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
@@ -16,16 +19,18 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
   onChange,
   disabled,
   loading,
+  style,
 }) => {
   return (
     <Select
       showSearch
-      placeholder="Search or select an organization"
+      placeholder="All Organizations"
       value={value}
       onChange={onChange}
       disabled={disabled}
       loading={loading}
       allowClear
+      style={{ minWidth: 280, ...style }}
       filterOption={(input, option) => {
         if (!option) return false;
         const org = organizations?.find((o) => o.organization_id === option.key);
@@ -37,12 +42,11 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
 
         return orgAlias.includes(searchTerm) || orgId.includes(searchTerm);
       }}
-      optionFilterProp="children"
     >
       {organizations?.map((org) => (
         <Select.Option key={org.organization_id} value={org.organization_id}>
           <span className="font-medium">{org.organization_alias}</span>{" "}
-          <span className="text-gray-500">({org.organization_id})</span>
+          <Text type="secondary">({org.organization_id})</Text>
         </Select.Option>
       ))}
     </Select>
diff --git a/ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx b/ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx
new file mode 100644
index 0000000000..9e90f77584
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx
@@ -0,0 +1,12 @@
+import { Spin } from "antd";
+import { LoadingOutlined } from "@ant-design/icons";
+
+interface AntDLoadingSpinnerProps {
+  size?: "small" | "default" | "large";
+  fontSize?: number;
+}
+
+export function AntDLoadingSpinner({ size, fontSize }: AntDLoadingSpinnerProps) {
+  const indicator = <LoadingOutlined style={fontSize ? { fontSize } : undefined} spin />;
+  return <Spin indicator={indicator} size={size} />;
+}

From ad43a35d762d62cd2cb9a18a0f1fbdef4f3fb67c Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <rcrabbe@berkeley.edu>
Date: Thu, 19 Mar 2026 22:50:19 -0700
Subject: [PATCH 15/31] feat: add control plane for multi-proxy worker
 management
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds a control plane capability that enables a central admin instance
to manage multiple regional worker proxies from a single UI.

Backend:
- Worker registry loaded from YAML config (worker_id, name, url)
- /.well-known/litellm-ui-config exposes is_control_plane and workers list
- /v3/login + /v3/login/exchange: opaque code exchange for cross-origin
  username/password auth (JWT never in URL/logs, single-use 60s TTL)
- SSO cookie handoff with return_to → opaque code → exchange
- _validate_return_to: full origin validation (scheme+hostname+port)
- Startup warning when control_plane_url set without Redis
- Both /v3 endpoints gated behind control_plane_url config

Frontend:
- Worker selector dropdown on login page (gated behind is_control_plane)
- Cross-origin SSO code exchange handling on callback
- switchToWorkerUrl: localStorage-persisted worker URL for API calls
- useWorker hook: shared worker state management
- WorkerDropdown in navbar for switching workers
- Logout/switch clears worker state from localStorage

Tests:
- 7 tests for /v3/login + /v3/login/exchange
- 10 tests for _validate_return_to
- 2 tests for control plane discovery endpoint
---
 .../ui_discovery_endpoints.py                 |   7 +
 litellm/proxy/management_endpoints/ui_sso.py  |  85 ++++++-
 litellm/proxy/proxy_server.py                 | 187 +++++++++++++++-
 .../types/proxy/control_plane_endpoints.py    |  14 ++
 .../ui_discovery_endpoints.py                 |   6 +-
 .../test_ui_discovery_endpoints.py            |  57 ++++-
 .../proxy/management_endpoints/test_ui_sso.py |  98 +++++++-
 tests/test_litellm/proxy/test_proxy_server.py | 211 ++++++++++++++++++
 .../app/(dashboard)/hooks/login/useLogin.ts   |   4 +-
 .../hooks/uiConfig/useUIConfig.test.ts        |   4 +
 .../src/app/login/LoginPage.test.tsx          |  15 +-
 .../src/app/login/LoginPage.tsx               | 125 +++++++++--
 .../Navbar/WorkerDropdown/WorkerDropdown.tsx  |  38 ++++
 .../src/components/navbar.tsx                 |  13 ++
 .../src/components/networking.tsx             | 126 ++++++++++-
 ui/litellm-dashboard/src/hooks/useWorker.ts   |  65 ++++++
 16 files changed, 1025 insertions(+), 30 deletions(-)
 create mode 100644 litellm/types/proxy/control_plane_endpoints.py
 create mode 100644 ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx
 create mode 100644 ui/litellm-dashboard/src/hooks/useWorker.ts

diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
index 2a38ceffba..233df5c6c5 100644
--- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
+++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
@@ -27,10 +27,17 @@ async def get_ui_config():
     admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"
 
     sso_configured = _has_user_setup_sso()
+
+    from litellm.proxy.proxy_server import proxy_config
+
+    is_control_plane = len(proxy_config.worker_registry) > 0
+
     return UiDiscoveryEndpoints(
         server_root_path=get_server_root_path(),
         proxy_base_url=get_proxy_base_url(),
         auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso,
         admin_ui_disabled=admin_ui_disabled,
         sso_configured=sso_configured,
+        is_control_plane=is_control_plane,
+        workers=proxy_config.worker_registry if is_control_plane else [],
     )
diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index daf1d6f131..8634a8bb33 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -16,6 +16,7 @@ import os
 import secrets
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
+from urllib.parse import urlencode, urlparse
 
 if TYPE_CHECKING:
     import httpx
@@ -301,6 +302,7 @@ async def google_login(
     source: Optional[str] = None,
     key: Optional[str] = None,
     existing_key: Optional[str] = None,
+    return_to: Optional[str] = None,
 ):  # noqa: PLR0915
     """
     Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
@@ -394,13 +396,23 @@ async def google_login(
         is True
     ):
         verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}")
-        return await SSOAuthenticationHandler.get_sso_login_redirect(
+        sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect(
             redirect_url=redirect_url,
             microsoft_client_id=microsoft_client_id,
             google_client_id=google_client_id,
             generic_client_id=generic_client_id,
             state=cli_state,
         )
+        if return_to is not None and sso_redirect is not None:
+            SSOAuthenticationHandler._validate_return_to(return_to)
+            sso_redirect.set_cookie(
+                key="litellm_cp_return_to",
+                value=return_to,
+                max_age=600,
+                httponly=True,
+                samesite="lax",
+            )
+        return sso_redirect
     elif ui_username is not None:
         # No Google, Microsoft SSO
         # Use UI Credentials set in .env
@@ -1312,12 +1324,17 @@ async def auth_callback(request: Request, state: Optional[str] = None):  # noqa:
             request=request, key=key_id, existing_key=existing_key, result=result
         )
 
+    # Control-plane cross-origin: read return_to from cookie.
+    # Starlette's cookie_parser already handles RFC 2109 unquoting.
+    cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to")
+
     return await SSOAuthenticationHandler.get_redirect_response_from_openid(
         result=result,
         request=request,
         received_response=received_response,
         generic_client_id=generic_client_id,
         ui_access_mode=ui_access_mode,
+        return_to=cp_return_to,
     )
 
 
@@ -1760,6 +1777,38 @@ class SSOAuthenticationHandler:
     Handler for SSO Authentication across all SSO providers
     """
 
+    @staticmethod
+    def _validate_return_to(return_to: str) -> None:
+        """
+        Validate that return_to matches the configured control_plane_url origin.
+
+        Raises HTTPException(400) if:
+        - control_plane_url is not configured in general_settings
+        - return_to origin does not match control_plane_url origin
+        """
+        from litellm.proxy.proxy_server import general_settings
+
+        control_plane_url = general_settings.get("control_plane_url")
+        if control_plane_url is None:
+            raise HTTPException(
+                status_code=400,
+                detail="return_to is not allowed: control_plane_url is not configured",
+            )
+
+        def _origin(url: str) -> tuple:
+            parsed = urlparse(url)
+            scheme = (parsed.scheme or "").lower()
+            hostname = (parsed.hostname or "").lower()
+            default_port = 443 if scheme == "https" else 80
+            port = parsed.port if parsed.port is not None else default_port
+            return (scheme, hostname, port)
+
+        if _origin(return_to) != _origin(control_plane_url):
+            raise HTTPException(
+                status_code=400,
+                detail="return_to does not match the configured control_plane_url",
+            )
+
     @staticmethod
     async def get_sso_login_redirect(
         redirect_url: str,
@@ -2358,6 +2407,7 @@ class SSOAuthenticationHandler:
         received_response: Optional[dict] = None,
         generic_client_id: Optional[str] = None,
         ui_access_mode: Optional[Dict] = None,
+        return_to: Optional[str] = None,
     ) -> RedirectResponse:
         import jwt
 
@@ -2367,6 +2417,7 @@ class SSOAuthenticationHandler:
             master_key,
             premium_user,
             proxy_logging_obj,
+            redis_usage_cache,
             user_api_key_cache,
             user_custom_sso,
         )
@@ -2534,6 +2585,38 @@ class SSOAuthenticationHandler:
             master_key or "",
             algorithm="HS256",
         )
+
+        # Control-plane cross-origin: store JWT behind a single-use opaque
+        # code (60s TTL) so the token never appears in browser history / logs.
+        # The control plane redeems it via POST /v3/login/exchange.
+        if return_to is not None:
+            SSOAuthenticationHandler._validate_return_to(return_to)
+
+            code = secrets.token_urlsafe(32)
+            cache_key = f"login_code:{code}"
+            cache_value = {"token": jwt_token, "redirect_url": return_to}
+            if redis_usage_cache is not None:
+                await redis_usage_cache.async_set_cache(
+                    key=cache_key, value=cache_value, ttl=60
+                )
+            else:
+                await user_api_key_cache.async_set_cache(
+                    key=cache_key, value=cache_value, ttl=60
+                )
+
+            separator = "&" if "?" in return_to else "?"
+            redirect_url = (
+                return_to
+                + separator
+                + urlencode({"login": "success", "code": code})
+            )
+            verbose_proxy_logger.info(
+                "Cross-origin SSO: redirecting to control plane with login code"
+            )
+            redirect_response = RedirectResponse(url=redirect_url, status_code=303)
+            redirect_response.delete_cookie("litellm_cp_return_to")
+            return redirect_response
+
         if user_id is not None and isinstance(user_id, str):
             litellm_dashboard_ui += "?login=success"
         verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 9c29927c5c..e982c934aa 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -541,6 +541,7 @@ from litellm.types.llms.anthropic import (
     AnthropicResponseUsageBlock,
 )
 from litellm.types.llms.openai import HttpxBinaryResponseContent
+from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
 from litellm.types.proxy.management_endpoints.model_management_endpoints import (
     ModelGroupInfoProxy,
 )
@@ -1546,6 +1547,7 @@ user_custom_key_generate = None
 # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload.
 # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'.
 _pkce_no_redis_warning_emitted: bool = False
+_cp_no_redis_warning_emitted: bool = False
 user_custom_sso = None
 user_custom_ui_sso_sign_in_handler = None
 use_background_health_checks = None
@@ -2295,6 +2297,7 @@ class ProxyConfig:
         self.config: Dict[str, Any] = {}
         self._last_semantic_filter_config: Optional[Dict[str, Any]] = None
         self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None
+        self.worker_registry: List["WorkerRegistryEntry"] = []
 
     def is_yaml(self, config_file_path: str) -> bool:
         if not os.path.isfile(config_file_path):
@@ -3095,6 +3098,21 @@ class ProxyConfig:
                         "Set PKCE_STRICT_CACHE_MISS=true to fail fast with a 401 on cache misses "
                         "instead of continuing without a code_verifier."
                     )
+
+            ### CONTROL PLANE CODE-EXCHANGE PREREQUISITE CHECK ###
+            cp_url = general_settings.get("control_plane_url")
+            if cp_url and redis_usage_cache is None:
+                global _cp_no_redis_warning_emitted
+                if not _cp_no_redis_warning_emitted:
+                    _cp_no_redis_warning_emitted = True
+                    verbose_proxy_logger.warning(
+                        "control_plane_url is configured but Redis is not configured for LiteLLM caching. "
+                        "Login codes (SSO and /v3/login) will not be shared across instances — "
+                        "the /v3/login/exchange call may land on a different pod and fail with 401. "
+                        "Configure Redis via the 'cache' section in your proxy config, "
+                        "or ensure sticky sessions for single-instance deployments."
+                    )
+
             ### STORE MODEL IN DB ### feature flag for `/model/new`
             store_model_in_db = general_settings.get("store_model_in_db", False)
             if store_model_in_db is None:
@@ -3385,7 +3403,15 @@ class ProxyConfig:
             litellm.vector_store_registry.load_vector_stores_from_config(
                 vector_store_registry_config
             )
-        pass
+
+        ## WORKER REGISTRY (Control Plane)
+        worker_registry_config = config.get("worker_registry", None)
+        if worker_registry_config:
+            self.worker_registry = [
+                WorkerRegistryEntry(**e) for e in worker_registry_config
+            ]
+        else:
+            self.worker_registry = []
 
     async def _init_policy_engine(
         self,
@@ -11095,6 +11121,165 @@ async def login_v2(request: Request):  # noqa: PLR0915
             )
 
 
+@router.post(
+    "/v3/login", include_in_schema=False
+)  # control-plane login — always returns token in body for cross-origin use
+async def login_v3(request: Request):  # noqa: PLR0915
+    global premium_user, general_settings, master_key
+    from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
+    from litellm.proxy.utils import get_custom_url
+
+    try:
+        if not general_settings.get("control_plane_url"):
+            raise ProxyException(
+                message="/v3/login is only available on workers with control_plane_url configured",
+                type=ProxyErrorTypes.not_found_error,
+                param="control_plane_url",
+                code=status.HTTP_404_NOT_FOUND,
+            )
+
+        body = await request.json()
+        username = str(body.get("username"))
+        password = str(body.get("password"))
+
+        login_result = await authenticate_user(
+            username=username,
+            password=password,
+            master_key=master_key,
+            prisma_client=prisma_client,
+        )
+
+        returned_ui_token_object = create_ui_token_object(
+            login_result=login_result,
+            general_settings=general_settings,
+            premium_user=premium_user,
+        )
+
+        import jwt
+
+        jwt_token = jwt.encode(
+            cast(dict, returned_ui_token_object),
+            cast(str, master_key),
+            algorithm="HS256",
+        )
+
+        litellm_dashboard_ui = get_custom_url(str(request.base_url))
+        if litellm_dashboard_ui.endswith("/"):
+            litellm_dashboard_ui += "ui/"
+        else:
+            litellm_dashboard_ui += "/ui/"
+        litellm_dashboard_ui += "?login=success"
+
+        # Store JWT behind a single-use opaque code (60s TTL)
+        code = secrets.token_urlsafe(32)
+        cache_key = f"login_code:{code}"
+        cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui}
+        if redis_usage_cache is not None:
+            await redis_usage_cache.async_set_cache(
+                key=cache_key, value=cache_value, ttl=60
+            )
+        else:
+            await user_api_key_cache.async_set_cache(
+                key=cache_key, value=cache_value, ttl=60
+            )
+
+        return JSONResponse(
+            content={"code": code, "expires_in": 60},
+            status_code=status.HTTP_200_OK,
+        )
+    except Exception as e:
+        verbose_proxy_logger.exception(
+            "litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format(
+                str(e)
+            )
+        )
+        if isinstance(e, ProxyException):
+            raise e
+        elif isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", str(e)),
+                type=ProxyErrorTypes.auth_error,
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
+            )
+        else:
+            error_msg = f"{str(e)}"
+            raise ProxyException(
+                message=error_msg,
+                type=ProxyErrorTypes.auth_error,
+                param="None",
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
+
+
+@router.post(
+    "/v3/login/exchange", include_in_schema=False
+)  # exchange single-use opaque code for JWT
+async def login_v3_exchange(request: Request):
+    try:
+        if not general_settings.get("control_plane_url"):
+            raise ProxyException(
+                message="/v3/login/exchange is only available on workers with control_plane_url configured",
+                type=ProxyErrorTypes.not_found_error,
+                param="control_plane_url",
+                code=status.HTTP_404_NOT_FOUND,
+            )
+
+        body = await request.json()
+        code = body.get("code")
+        if not code:
+            raise ProxyException(
+                message="Missing 'code' parameter",
+                type=ProxyErrorTypes.auth_error,
+                param="code",
+                code=status.HTTP_400_BAD_REQUEST,
+            )
+
+        cache_key = f"login_code:{code}"
+        if redis_usage_cache is not None:
+            cached_data = await redis_usage_cache.async_get_cache(key=cache_key)
+        else:
+            cached_data = await user_api_key_cache.async_get_cache(key=cache_key)
+
+        if not cached_data or not isinstance(cached_data, dict):
+            raise ProxyException(
+                message="Invalid or expired login code",
+                type=ProxyErrorTypes.auth_error,
+                param="code",
+                code=status.HTTP_401_UNAUTHORIZED,
+            )
+
+        # Single-use: delete immediately
+        if redis_usage_cache is not None:
+            await redis_usage_cache.async_delete_cache(key=cache_key)
+        else:
+            await user_api_key_cache.async_delete_cache(key=cache_key)
+
+        json_response = JSONResponse(
+            content={
+                "token": cached_data["token"],
+                "redirect_url": cached_data["redirect_url"],
+            },
+            status_code=status.HTTP_200_OK,
+        )
+        json_response.set_cookie(key="token", value=cached_data["token"])
+        return json_response
+    except ProxyException:
+        raise
+    except Exception as e:
+        verbose_proxy_logger.exception(
+            "litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format(
+                str(e)
+            )
+        )
+        raise ProxyException(
+            message=str(e),
+            type=ProxyErrorTypes.auth_error,
+            param="None",
+            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+
+
 @app.get("/onboarding/get_token", include_in_schema=False)
 async def onboarding(invite_link: str, request: Request):
     """
diff --git a/litellm/types/proxy/control_plane_endpoints.py b/litellm/types/proxy/control_plane_endpoints.py
new file mode 100644
index 0000000000..8bf4c44b20
--- /dev/null
+++ b/litellm/types/proxy/control_plane_endpoints.py
@@ -0,0 +1,14 @@
+from pydantic import BaseModel, field_validator
+
+
+class WorkerRegistryEntry(BaseModel):
+    worker_id: str
+    name: str
+    url: str
+
+    @field_validator("url")
+    @classmethod
+    def url_must_be_http(cls, v: str) -> str:
+        if not v.startswith(("http://", "https://")):
+            raise ValueError("Worker URL must start with http:// or https://")
+        return v
diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py
index 4a4cdaa2ba..46cd3f49f1 100644
--- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py
+++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py
@@ -1,7 +1,9 @@
-from typing import Optional
+from typing import List, Optional
 
 from pydantic import BaseModel
 
+from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
+
 
 class UiDiscoveryEndpoints(BaseModel):
     server_root_path: str
@@ -9,3 +11,5 @@ class UiDiscoveryEndpoints(BaseModel):
     auto_redirect_to_sso: bool
     admin_ui_disabled: bool
     sso_configured: bool
+    is_control_plane: bool = False
+    workers: List[WorkerRegistryEntry] = []
diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py
index f15960a607..54a127f435 100644
--- a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py
+++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py
@@ -1,6 +1,6 @@
 import os
 import sys
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from fastapi import FastAPI
@@ -11,6 +11,7 @@ sys.path.insert(
 )
 
 from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import router
+from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
 
 
 def test_ui_discovery_endpoints_with_defaults():
@@ -245,9 +246,9 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled():
          patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
          patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
          patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
-        
+
         response = client.get("/.well-known/litellm-ui-config")
-        
+
         assert response.status_code == 200
         data = response.json()
         assert data["server_root_path"] == "/"
@@ -256,3 +257,53 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled():
         assert data["admin_ui_disabled"] is False
         assert data["sso_configured"] is False
 
+
+def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured():
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    mock_config = MagicMock()
+    mock_config.worker_registry = [
+        WorkerRegistryEntry(
+            worker_id="team-a", name="Team A", url="https://worker-1:4001"
+        ),
+    ]
+
+    with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
+         patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
+         patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
+         patch("litellm.proxy.proxy_server.proxy_config", mock_config), \
+         patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
+
+        response = client.get("/.well-known/litellm-ui-config")
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["is_control_plane"] is True
+        assert len(data["workers"]) == 1
+        assert data["workers"][0]["worker_id"] == "team-a"
+        assert data["workers"][0]["name"] == "Team A"
+        assert data["workers"][0]["url"] == "https://worker-1:4001"
+
+
+def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers():
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+
+    with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
+         patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
+         patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
+         patch("litellm.proxy.proxy_server.proxy_config", mock_config), \
+         patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
+
+        response = client.get("/.well-known/litellm-ui-config")
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["is_control_plane"] is False
+        assert data["workers"] == []
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
index d43b2c4ba0..fc9c37b7f8 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
-from fastapi import Request
+from fastapi import HTTPException, Request
 
 from litellm._uuid import uuid
 
@@ -5160,3 +5160,99 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch):
     assert result.extra_fields["missing_field"] is None
     assert result.extra_fields["another_missing"] is None
 
+
+class TestValidateReturnTo:
+    """Tests for SSOAuthenticationHandler._validate_return_to"""
+
+    def test_rejects_when_no_control_plane_url_configured(self, monkeypatch):
+        """return_to should be rejected if control_plane_url is not in general_settings."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings", {}
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
+        assert exc_info.value.status_code == 400
+        assert "not configured" in exc_info.value.detail
+
+    def test_allows_matching_origin(self, monkeypatch):
+        """return_to matching the configured control_plane_url origin should pass."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        # Should not raise
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models")
+
+    def test_allows_matching_origin_with_trailing_slash(self, monkeypatch):
+        """Trailing slash on control_plane_url should not affect origin comparison."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com/"},
+        )
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
+
+    def test_rejects_prefix_attack(self, monkeypatch):
+        """return_to like cp.example.com.evil.com must be rejected (not just prefix match)."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal")
+        assert exc_info.value.status_code == 400
+
+    def test_rejects_different_origin(self, monkeypatch):
+        """return_to pointing to a completely different domain should be rejected."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://evil.com/phish")
+        assert exc_info.value.status_code == 400
+
+    def test_case_insensitive_hostname(self, monkeypatch):
+        """Hostname comparison should be case-insensitive per RFC 3986."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://CP.Example.COM"},
+        )
+        # Should not raise
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
+
+    def test_rejects_scheme_mismatch(self, monkeypatch):
+        """http:// must be rejected when control_plane_url uses https://."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("http://cp.example.com/ui")
+        assert exc_info.value.status_code == 400
+
+    def test_rejects_port_mismatch(self, monkeypatch):
+        """Non-default port must be rejected."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui")
+        assert exc_info.value.status_code == 400
+
+    def test_allows_explicit_default_port(self, monkeypatch):
+        """https://host:443 should match https://host (default port normalisation)."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com:443/ui")
+
+    def test_allows_matching_custom_port(self, monkeypatch):
+        """Both sides on the same custom port should match."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com:3000"},
+        )
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com:3000/ui")
+
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 112a06b173..bd6162f225 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -236,6 +236,217 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch):
     assert isinstance(data["error"], dict)
 
 
+def test_login_v3_rejected_without_control_plane_url(monkeypatch):
+    """v3/login returns 404 when control_plane_url is not configured."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+
+    assert response.status_code == 404
+    assert "control_plane_url" in response.json()["error"]["message"]
+
+
+def test_login_v3_returns_code(monkeypatch):
+    """v3/login returns an opaque code, not the JWT directly."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        AsyncMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.create_ui_token_object",
+        MagicMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
+    monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
+    monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert "code" in data
+    assert data["expires_in"] == 60
+    assert "token" not in data
+
+
+def test_login_v3_exchange_happy_path(monkeypatch):
+    """Full flow: v3/login returns code, v3/login/exchange redeems it for JWT."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        AsyncMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.create_ui_token_object",
+        MagicMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
+    monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
+    monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
+
+    client = TestClient(app)
+
+    # Step 1: login — get code
+    login_response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+    assert login_response.status_code == 200
+    code = login_response.json()["code"]
+
+    # Step 2: exchange — get JWT
+    exchange_response = client.post(
+        "/v3/login/exchange",
+        json={"code": code},
+    )
+    assert exchange_response.status_code == 200
+    exchange_data = exchange_response.json()
+    assert exchange_data["token"] == "signed-token"
+    assert "redirect_url" in exchange_data
+    assert exchange_response.cookies.get("token") == "signed-token"
+
+
+def test_login_v3_exchange_single_use(monkeypatch):
+    """Code can only be redeemed once."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        AsyncMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.create_ui_token_object",
+        MagicMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
+    monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
+    monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
+
+    client = TestClient(app)
+
+    login_response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+    code = login_response.json()["code"]
+
+    # First exchange succeeds
+    first = client.post("/v3/login/exchange", json={"code": code})
+    assert first.status_code == 200
+
+    # Second exchange fails
+    second = client.post("/v3/login/exchange", json={"code": code})
+    assert second.status_code == 401
+
+
+def test_login_v3_exchange_invalid_code(monkeypatch):
+    """Random code returns 401."""
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login/exchange",
+        json={"code": "nonexistent-code"},
+    )
+    assert response.status_code == 401
+
+
+def test_login_v3_exchange_rejected_without_control_plane_url(monkeypatch):
+    """v3/login/exchange returns 404 when control_plane_url is not configured."""
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login/exchange",
+        json={"code": "some-code"},
+    )
+
+    assert response.status_code == 404
+    assert "control_plane_url" in response.json()["error"]["message"]
+
+
+def test_login_v3_returns_json_on_proxy_exception(monkeypatch):
+    """Test that /v3/login returns JSON error when ProxyException is raised"""
+    from litellm.proxy._types import ProxyErrorTypes, ProxyException
+
+    mock_prisma_client = MagicMock()
+    mock_authenticate_user = AsyncMock(
+        side_effect=ProxyException(
+            message="Invalid credentials",
+            type=ProxyErrorTypes.auth_error,
+            param="password",
+            code=401,
+        )
+    )
+
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        mock_authenticate_user,
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "wrong"},
+    )
+
+    assert response.status_code == 401
+    assert response.headers["content-type"] == "application/json"
+    data = response.json()
+    assert "error" in data
+    assert data["error"]["message"] == "Invalid credentials"
+    assert data["error"]["type"] == "auth_error"
+
+
 def test_fallback_login_has_no_deprecation_banner(client_no_auth):
     response = client_no_auth.get("/fallback/login")
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts
index a15b4a06d1..be53b1c80a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts
@@ -3,8 +3,8 @@ import { loginCall, LoginRequest } from "@/components/networking";
 
 export const useLogin = () => {
   return useMutation({
-    mutationFn: async ({ username, password }: LoginRequest) => {
-      const result = await loginCall(username, password);
+    mutationFn: async ({ username, password, useV3 }: LoginRequest) => {
+      const result = await loginCall(username, password, useV3);
       return result;
     },
   });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts
index aba5dddf13..b05bae1e81 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts
@@ -28,6 +28,8 @@ const mockUIConfig: LiteLLMWellKnownUiConfig = {
   proxy_base_url: "https://proxy.example.com",
   auto_redirect_to_sso: true,
   admin_ui_disabled: false,
+  is_control_plane: false,
+  workers: [],
 };
 
 describe("useUIConfig", () => {
@@ -102,6 +104,8 @@ describe("useUIConfig", () => {
       auto_redirect_to_sso: false,
       sso_configured: false,
       admin_ui_disabled: true,
+      is_control_plane: false,
+      workers: [],
     };
 
     // Mock successful API call with different data
diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx
index ad2dde2da8..866b7d0f17 100644
--- a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx
+++ b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx
@@ -41,6 +41,17 @@ vi.mock("@/app/(dashboard)/hooks/login/useLogin", () => ({
   })),
 }));
 
+vi.mock("@/hooks/useWorker", () => ({
+  useWorker: vi.fn(() => ({
+    isControlPlane: false,
+    workers: [],
+    selectedWorkerId: null,
+    selectedWorker: null,
+    selectWorker: vi.fn(),
+    disconnectFromWorker: vi.fn(),
+  })),
+}));
+
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import { getCookie } from "@/utils/cookieUtils";
 import { isJwtExpired } from "@/utils/jwtUtils";
@@ -108,7 +119,7 @@ describe("LoginPage", () => {
     );
 
     await waitFor(() => {
-      expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui");
+      expect(mockReplace).toHaveBeenCalledWith("/ui");
     });
   });
 
@@ -189,7 +200,7 @@ describe("LoginPage", () => {
     );
 
     await waitFor(() => {
-      expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui");
+      expect(mockReplace).toHaveBeenCalledWith("/ui");
     });
 
     expect(mockPush).not.toHaveBeenCalled();
diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx
index d54140c5a2..5a9d420456 100644
--- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx
+++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx
@@ -3,14 +3,15 @@
 import { useLogin } from "@/app/(dashboard)/hooks/login/useLogin";
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import LoadingScreen from "@/components/common_components/LoadingScreen";
-import { getProxyBaseUrl } from "@/components/networking";
-import { getCookie } from "@/utils/cookieUtils";
+import { exchangeLoginCode, getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking";
+import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
 import { isJwtExpired } from "@/utils/jwtUtils";
 import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils";
-import { InfoCircleOutlined } from "@ant-design/icons";
-import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd";
+import { InfoCircleOutlined, CloudServerOutlined } from "@ant-design/icons";
+import { Alert, Button, Card, Form, Input, Popover, Select, Space, Typography } from "antd";
 import { useRouter } from "next/navigation";
 import { useEffect, useState } from "react";
+import { useWorker } from "@/hooks/useWorker";
 
 function LoginPageContent() {
   const [username, setUsername] = useState("");
@@ -19,6 +20,17 @@ function LoginPageContent() {
   const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig();
   const loginMutation = useLogin();
   const router = useRouter();
+  const { workers, selectWorker } = useWorker();
+  const [selectedWorkerId, setSelectedWorkerId] = useState<string | null>(null);
+
+  // Pre-select worker from URL param (e.g. /ui/login?worker=team-b)
+  useEffect(() => {
+    const params = new URLSearchParams(window.location.search);
+    const workerParam = params.get("worker");
+    if (workerParam) {
+      setSelectedWorkerId(workerParam);
+    }
+  }, []);
 
   useEffect(() => {
     if (isConfigLoading) {
@@ -31,6 +43,44 @@ function LoginPageContent() {
       return;
     }
 
+    // Cross-origin SSO: worker redirected back with a single-use code.
+    // Exchange it for the JWT via the worker's /v3/login/exchange endpoint.
+    const params = new URLSearchParams(window.location.search);
+    const ssoCode = params.get("code");
+    if (ssoCode) {
+      const workerUrl = localStorage.getItem("litellm_worker_url");
+      exchangeLoginCode(ssoCode, workerUrl).then(() => {
+        params.delete("code");
+        const cleanSearch = params.toString();
+        window.history.replaceState(null, "", window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""));
+        router.replace("/ui/?login=success");
+      });
+      return;
+    }
+
+    // Backwards compat: handle direct token in URL (legacy flow)
+    const urlToken = params.get("token");
+    if (urlToken && !isJwtExpired(urlToken)) {
+      document.cookie = `token=${urlToken}; path=/; SameSite=Lax`;
+      params.delete("token");
+      const cleanSearch = params.toString();
+      window.history.replaceState(
+        null,
+        "",
+        window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""),
+      );
+      router.replace("/ui/?login=success");
+      return;
+    }
+
+    // If switching workers on a control plane, clear the old token and show login
+    const switchingWorker = params.has("worker");
+    if (switchingWorker && uiConfig?.is_control_plane) {
+      clearTokenCookies();
+      setIsLoading(false);
+      return;
+    }
+
     const rawToken = getCookie("token");
     if (rawToken && !isJwtExpired(rawToken)) {
       // User already logged in - redirect to return URL or default
@@ -38,7 +88,7 @@ function LoginPageContent() {
       if (returnUrl) {
         router.replace(returnUrl);
       } else {
-        router.replace(`${getProxyBaseUrl()}/ui`);
+        router.replace("/ui");
       }
       return;
     }
@@ -58,16 +108,35 @@ function LoginPageContent() {
   }, [isConfigLoading, router, uiConfig]);
 
   const handleSubmit = () => {
+    // If a worker is selected, point proxyBaseUrl at it before login
+    const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
+    if (selectedWorker) {
+      switchToWorkerUrl(selectedWorker.url);
+    }
+
     loginMutation.mutate(
-      { username, password },
+      { username, password, useV3: !!selectedWorker },
       {
         onSuccess: (data) => {
-          // Check if we have a return URL to use instead of the default redirect
-          const returnUrl = consumeReturnUrl();
-          if (returnUrl) {
-            router.push(returnUrl);
+          // Update the worker context with the selected worker
+          if (selectedWorker) {
+            selectWorker(selectedWorker.worker_id);
+            // Stay on the CP's UI — proxyBaseUrl already points at the worker
+            router.push("/ui/?login=success");
           } else {
-            router.push(data.redirect_url);
+            // Normal (non-control-plane) login — follow the server's redirect
+            const returnUrl = consumeReturnUrl();
+            if (returnUrl) {
+              router.push(returnUrl);
+            } else {
+              router.push(data.redirect_url);
+            }
+          }
+        },
+        onError: () => {
+          // Reset proxyBaseUrl on login failure
+          if (selectedWorker) {
+            switchToWorkerUrl(null);
           }
         },
       },
@@ -154,6 +223,22 @@ function LoginPageContent() {
           {error && <Alert message={error} type="error" showIcon />}
 
           <Form onFinish={handleSubmit} layout="vertical" requiredMark={true}>
+            {uiConfig?.is_control_plane && workers.length > 0 && (
+              <Form.Item label="Worker" style={{ marginBottom: 16 }}>
+                <Select
+                  value={selectedWorkerId || undefined}
+                  onChange={(value) => setSelectedWorkerId(value)}
+                  placeholder="Choose a worker to connect to"
+                  size="large"
+                  suffixIcon={<CloudServerOutlined />}
+                  options={workers.map((w) => ({
+                    label: w.name,
+                    value: w.worker_id,
+                  }))}
+                />
+              </Form.Item>
+            )}
+
             <Form.Item
               label="Username"
               name="username"
@@ -209,10 +294,20 @@ function LoginPageContent() {
                 </Popover>
               ) : (
                 <Button
-                  disabled={isLoginLoading}
-                  onClick={() =>
-                    router.push(`${getProxyBaseUrl()}/sso/key/generate`)
-                  }
+                  disabled={isLoginLoading || (!!selectedWorkerId && workers.length === 0)}
+                  onClick={() => {
+                    const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
+                    if (selectedWorker) {
+                      // Store worker selection so useWorker hook restores it after redirect
+                      localStorage.setItem("litellm_selected_worker_id", selectedWorkerId!);
+                      switchToWorkerUrl(selectedWorker.url);
+                    }
+                    // SSO on the worker (or this instance if no worker), always
+                    // include return_to so the callback redirects back here
+                    const ssoBase = selectedWorker?.url ?? getProxyBaseUrl();
+                    const returnTo = encodeURIComponent(window.location.origin + "/ui/login");
+                    router.push(`${ssoBase}/sso/key/generate?return_to=${returnTo}`);
+                  }}
                   block
                   size="large"
                 >
diff --git a/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx b/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx
new file mode 100644
index 0000000000..55309a11c1
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx
@@ -0,0 +1,38 @@
+"use client";
+
+import React from "react";
+import { Select } from "antd";
+import { CloudServerOutlined } from "@ant-design/icons";
+import { useWorker } from "@/hooks/useWorker";
+
+interface WorkerDropdownProps {
+  onWorkerSwitch: (workerId: string) => void;
+}
+
+const WorkerDropdown: React.FC<WorkerDropdownProps> = ({ onWorkerSwitch }) => {
+  const { isControlPlane, selectedWorker, workers } = useWorker();
+
+  if (!isControlPlane || !selectedWorker) return null;
+
+  return (
+    <Select
+      showSearch
+      filterOption={(input, option) =>
+        (option?.label as string ?? "").toLowerCase().includes(input.toLowerCase())
+      }
+      value={selectedWorker.worker_id}
+      style={{ minWidth: 180 }}
+      suffixIcon={<CloudServerOutlined />}
+      options={workers.map((w) => ({
+        label: w.name,
+        value: w.worker_id,
+        disabled: w.worker_id === selectedWorker.worker_id,
+      }))}
+      onChange={(newWorkerId) => {
+        onWorkerSwitch(newWorkerId);
+      }}
+    />
+  );
+};
+
+export default WorkerDropdown;
diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx
index c46a3af5a6..96d6ce613b 100644
--- a/ui/litellm-dashboard/src/components/navbar.tsx
+++ b/ui/litellm-dashboard/src/components/navbar.tsx
@@ -4,6 +4,7 @@ import { getProxyBaseUrl } from "@/components/networking";
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import { useTheme } from "@/contexts/ThemeContext";
 import { clearTokenCookies } from "@/utils/cookieUtils";
+import { clearStoredReturnUrl } from "@/utils/returnUrlUtils";
 import { fetchProxySettings } from "@/utils/proxyUtils";
 import { MenuFoldOutlined, MenuUnfoldOutlined, MessageOutlined, MoonOutlined, SunOutlined } from "@ant-design/icons";
 import { Button, Switch, Tag } from "antd";
@@ -12,6 +13,7 @@ import React, { useEffect, useState } from "react";
 import { BlogDropdown } from "./Navbar/BlogDropdown/BlogDropdown";
 import { CommunityEngagementButtons } from "./Navbar/CommunityEngagementButtons/CommunityEngagementButtons";
 import UserDropdown from "./Navbar/UserDropdown/UserDropdown";
+import WorkerDropdown from "./Navbar/WorkerDropdown/WorkerDropdown";
 
 interface NavbarProps {
   userID: string | null;
@@ -77,9 +79,19 @@ const Navbar: React.FC<NavbarProps> = ({
 
   const handleLogout = () => {
     clearTokenCookies();
+    localStorage.removeItem("litellm_selected_worker_id");
+    localStorage.removeItem("litellm_worker_url");
     window.location.href = logoutUrl;
   };
 
+  const handleWorkerSwitch = (workerId: string) => {
+    clearTokenCookies();
+    clearStoredReturnUrl();
+    localStorage.removeItem("litellm_selected_worker_id");
+    localStorage.removeItem("litellm_worker_url");
+    window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`;
+  };
+
   return (
     <nav className="bg-white border-b border-gray-200 sticky top-0 z-10">
       <div className="w-full">
@@ -169,6 +181,7 @@ const Navbar: React.FC<NavbarProps> = ({
                 NEW
               </span>
             </a>
+            <WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />
             <CommunityEngagementButtons />
             {/* Dark mode is currently a work in progress. To test, you can change 'false' to 'true' below.
             Do not set this to true by default until all components are confirmed to support dark mode styles. */}
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index bb9a0c1e01..80170a6986 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -86,7 +86,23 @@ const defaultProxyBaseUrl =
     : null;
 const defaultServerRootPath = "/";
 export let serverRootPath = defaultServerRootPath;
-export let proxyBaseUrl = defaultProxyBaseUrl;
+const WORKER_URL_KEY = "litellm_worker_url";
+// If a worker URL is in localStorage, use it as the initial proxyBaseUrl.
+// This survives page navigation and the sessionStorage.clear() in user_dashboard.
+const _rawWorkerUrl =
+  typeof window !== "undefined" ? window.localStorage.getItem(WORKER_URL_KEY) : null;
+// Validate stored worker URL — reject non-HTTP schemes to prevent exfiltration
+const _initialWorkerUrl = (() => {
+  if (!_rawWorkerUrl) return null;
+  try {
+    const parsed = new URL(_rawWorkerUrl);
+    if (parsed.protocol === "http:" || parsed.protocol === "https:") return _rawWorkerUrl;
+  } catch { /* invalid URL */ }
+  // Invalid URL in storage — clear it
+  if (typeof window !== "undefined") window.localStorage.removeItem(WORKER_URL_KEY);
+  return null;
+})();
+export let proxyBaseUrl: string | null = _initialWorkerUrl ?? defaultProxyBaseUrl;
 if (isLocal != true) {
   console.log = function () { };
 }
@@ -102,6 +118,10 @@ const updateProxyBaseUrl = (serverRootPath: string, receivedProxyBaseUrl: string
   /**
    * Special function for updating the proxy base url. Should only be called by getUiConfig.
    */
+  // If a worker URL is in localStorage, don't let getUiConfig overwrite it
+  if (typeof window !== "undefined" && window.localStorage.getItem(WORKER_URL_KEY)) {
+    return;
+  }
   const browserLocation = getWindowLocation();
   const resolvedDefaultProxyBaseUrl =
     isLocal && process.env.NEXT_PUBLIC_USE_REWRITES !== "true"
@@ -137,6 +157,36 @@ export const getProxyBaseUrl = (): string => {
   return browserLocation?.origin ?? "";
 };
 
+/**
+ * Switch API calls to point at a worker (or back to the control plane).
+ * Persists to localStorage so it survives page navigation and the
+ * sessionStorage.clear() in user_dashboard. Also updates the module-level
+ * proxyBaseUrl so in-flight code in this JS execution sees the new value
+ * immediately.
+ */
+function isValidHttpUrl(url: string): boolean {
+  try {
+    const parsed = new URL(url);
+    return parsed.protocol === "http:" || parsed.protocol === "https:";
+  } catch {
+    return false;
+  }
+}
+
+export function switchToWorkerUrl(workerUrl: string | null): void {
+  if (workerUrl && !isValidHttpUrl(workerUrl)) {
+    return;
+  }
+  if (typeof window !== "undefined") {
+    if (workerUrl) {
+      window.localStorage.setItem(WORKER_URL_KEY, workerUrl);
+    } else {
+      window.localStorage.removeItem(WORKER_URL_KEY);
+    }
+  }
+  proxyBaseUrl = workerUrl ?? defaultProxyBaseUrl;
+}
+
 const HTTP_REQUEST = {
   GET: "GET",
   POST: "POST",
@@ -262,12 +312,20 @@ interface PublicModelHubInfo {
   useful_links: Record<string, string | { url: string; index: number }>;
 }
 
+export interface WorkerInfo {
+  worker_id: string;
+  name: string;
+  url: string;
+}
+
 export interface LiteLLMWellKnownUiConfig {
   server_root_path: string;
   proxy_base_url: string | null;
   auto_redirect_to_sso: boolean;
   admin_ui_disabled: boolean;
   sso_configured: boolean;
+  is_control_plane?: boolean;
+  workers?: WorkerInfo[];
 }
 
 export interface CredentialsResponse {
@@ -9030,15 +9088,20 @@ export const deriveErrorMessage = (errorData: any): string => {
 export interface LoginRequest {
   username: string;
   password: string;
+  useV3?: boolean;
 }
 
 interface LoginResponse {
   redirect_url: string;
+  token?: string;
+  code?: string;
+  expires_in?: number;
 }
 
-export const loginCall = async (username: string, password: string): Promise<LoginResponse> => {
+export const loginCall = async (username: string, password: string, useV3?: boolean): Promise<LoginResponse> => {
   const proxyBaseUrl = getProxyBaseUrl();
-  const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}/v2/login` : "/v2/login";
+  const loginPath = useV3 ? "/v3/login" : "/v2/login";
+  const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}${loginPath}` : loginPath;
 
   const body = JSON.stringify({
     username,
@@ -9060,10 +9123,65 @@ export const loginCall = async (username: string, password: string): Promise<Log
     throw new Error(errorMessage);
   }
 
-  const data = await response.json();
+  const data: LoginResponse = await response.json();
+
+  // v3 returns an opaque code — exchange it for the real JWT
+  if (useV3 && data.code) {
+    const exchangeUrl = proxyBaseUrl
+      ? `${proxyBaseUrl}/v3/login/exchange`
+      : "/v3/login/exchange";
+
+    const exchangeResponse = await fetch(exchangeUrl, {
+      method: "POST",
+      body: JSON.stringify({ code: data.code }),
+      credentials: "include",
+      headers: { "Content-Type": "application/json" },
+    });
+
+    if (!exchangeResponse.ok) {
+      const errorData = await exchangeResponse.json();
+      throw new Error(deriveErrorMessage(errorData));
+    }
+
+    const exchangeData: LoginResponse = await exchangeResponse.json();
+    if (exchangeData.token) {
+      document.cookie = `token=${exchangeData.token}; path=/; SameSite=Lax`;
+    }
+    return exchangeData;
+  }
+
+  // Backwards compatibility: v2 or old v3 returns token directly
+  if (data.token) {
+    document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
+  }
+
   return data;
 };
 
+/**
+ * Exchange a single-use login code for a JWT token.
+ * Used by the SSO callback when the worker redirects back with ?code=.
+ */
+export const exchangeLoginCode = async (code: string, workerBaseUrl?: string | null): Promise<string> => {
+  const base = workerBaseUrl || getProxyBaseUrl();
+  const response = await fetch(`${base}/v3/login/exchange`, {
+    method: "POST",
+    body: JSON.stringify({ code }),
+    headers: { "Content-Type": "application/json" },
+  });
+
+  if (!response.ok) {
+    const errorData = await response.json();
+    throw new Error(deriveErrorMessage(errorData));
+  }
+
+  const data = await response.json();
+  if (data.token) {
+    document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
+  }
+  return data.token;
+};
+
 export const getUiSettings = async () => {
   const proxyBaseUrl = getProxyBaseUrl();
   const url = proxyBaseUrl ? `${proxyBaseUrl}/get/ui_settings` : `/get/ui_settings`;
diff --git a/ui/litellm-dashboard/src/hooks/useWorker.ts b/ui/litellm-dashboard/src/hooks/useWorker.ts
new file mode 100644
index 0000000000..5afdd3d599
--- /dev/null
+++ b/ui/litellm-dashboard/src/hooks/useWorker.ts
@@ -0,0 +1,65 @@
+"use client";
+
+import { useState, useCallback, useEffect } from "react";
+import { switchToWorkerUrl, WorkerInfo } from "@/components/networking";
+import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
+
+const SELECTED_WORKER_KEY = "litellm_selected_worker_id";
+
+interface UseWorkerReturn {
+  isControlPlane: boolean;
+  workers: WorkerInfo[];
+  selectedWorkerId: string | null;
+  selectedWorker: WorkerInfo | null;
+  selectWorker: (workerId: string) => void;
+  disconnectFromWorker: () => void;
+}
+
+export const useWorker = (): UseWorkerReturn => {
+  const { data: uiConfig } = useUIConfig();
+  const isControlPlane = uiConfig?.is_control_plane ?? false;
+  const workers: WorkerInfo[] = uiConfig?.workers ?? [];
+
+  const [selectedWorkerId, setSelectedWorkerId] = useState<string | null>(() => {
+    if (typeof window === "undefined") return null;
+    return localStorage.getItem(SELECTED_WORKER_KEY);
+  });
+
+  // Once workers are loaded, restore proxyBaseUrl from the persisted selection
+  useEffect(() => {
+    if (!selectedWorkerId || workers.length === 0) return;
+    const worker = workers.find((w) => w.worker_id === selectedWorkerId);
+    if (worker) {
+      switchToWorkerUrl(worker.url);
+    }
+  }, [selectedWorkerId, workers]);
+
+  const selectedWorker =
+    workers.find((w) => w.worker_id === selectedWorkerId) ?? null;
+
+  const selectWorker = useCallback(
+    (workerId: string) => {
+      const worker = workers.find((w) => w.worker_id === workerId);
+      if (!worker) return;
+      setSelectedWorkerId(workerId);
+      localStorage.setItem(SELECTED_WORKER_KEY, workerId);
+      switchToWorkerUrl(worker.url);
+    },
+    [workers],
+  );
+
+  const disconnectFromWorker = useCallback(() => {
+    setSelectedWorkerId(null);
+    localStorage.removeItem(SELECTED_WORKER_KEY);
+    switchToWorkerUrl(null);
+  }, []);
+
+  return {
+    isControlPlane,
+    workers,
+    selectedWorkerId,
+    selectedWorker,
+    selectWorker,
+    disconnectFromWorker,
+  };
+};

From f4e5fb2b4ae4d388d62e66eaba64c33a162bf8d6 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 22:50:35 -0700
Subject: [PATCH 16/31] fix: exclude team_id from Prisma create payload in
 _prepare_mcp_server_data

team_id is a request-level field, not a DB column. Excluding it prevents
Prisma from rejecting the create call.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/proxy/_experimental/mcp_server/db.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py
index fbef33c32e..0bc102889b 100644
--- a/litellm/proxy/_experimental/mcp_server/db.py
+++ b/litellm/proxy/_experimental/mcp_server/db.py
@@ -40,8 +40,8 @@ def _prepare_mcp_server_data(
     """
     from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
-    # Convert model to dict
-    data_dict = data.model_dump(exclude_none=True)
+    # Convert model to dict, excluding fields not in the DB schema
+    data_dict = data.model_dump(exclude_none=True, exclude={"team_id"})
     # Ensure alias is always present in the dict (even if None)
     if "alias" not in data_dict:
         data_dict["alias"] = getattr(data, "alias", None)

From 6c06e6dc80fd7a2870a3f6fef57a8d651de1a82b Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 22:57:27 -0700
Subject: [PATCH 17/31] fix: link new object_permission_id back to team on
 first MCP server create

When a team has no object_permission_id yet, the auto-assign logic creates
an ObjectPermissionTable row but never linked it to the team. Now updates
the team's object_permission_id after creation.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../management_endpoints/mcp_management_endpoints.py     | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index 3400f90f48..522a533d43 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -1358,7 +1358,7 @@ if MCP_AVAILABLE:
                     set(existing_mcp_servers + [new_mcp_server.server_id])
                 )
 
-                await handle_update_object_permission_common(
+                new_permission_id = await handle_update_object_permission_common(
                     data_json={
                         "object_permission": {
                             "mcp_servers": updated_mcp_servers,
@@ -1367,6 +1367,13 @@ if MCP_AVAILABLE:
                     existing_object_permission_id=existing_permission_id,
                     prisma_client=prisma_client,
                 )
+
+                # If the team had no object_permission_id, link the new one
+                if existing_permission_id is None and new_permission_id is not None:
+                    await prisma_client.db.litellm_teamtable.update(
+                        where={"team_id": manager_team_id},
+                        data={"object_permission_id": new_permission_id},
+                    )
         except Exception as e:
             verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}")
             raise HTTPException(

From dfc4401e1c305f11c898e989991503122668bede Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 23:00:38 -0700
Subject: [PATCH 18/31] fix: use check_db_only=True in
 _assert_can_manage_team_mcp_server

Cached team objects may have stale object_permission data. Bypassing
cache ensures the server-in-team check uses fresh data.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/proxy/management_endpoints/mcp_management_endpoints.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index 522a533d43..a3b0b3ab28 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -179,7 +179,7 @@ if MCP_AVAILABLE:
             team_id=resolved_team_id,
             prisma_client=prisma_client,
             user_api_key_cache=user_api_key_cache,
-            check_db_only=False,
+            check_db_only=True,  # bypass cache to get fresh object_permission
         )
 
         if not _is_user_team_mcp_manager(user_api_key_dict, team_obj):

From e71410470267f64a00f61d7b75094a8628054b0e Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 23:10:44 -0700
Subject: [PATCH 19/31] fix: address Greptile review feedback

- P0: Separate try/except for auto-assign so server creation succeeds
  even if team permission update fails
- P1: Clean up team permission entry on MCP server delete
- P1: Add MCP_AVAILABLE skip guard to tests
- P2: Return team_obj from _assert_can_manage_team_mcp_server to
  eliminate redundant get_team_object call in create endpoint

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../mcp_management_endpoints.py               | 80 +++++++++++++------
 .../test_mcp_manager_role.py                  | 21 +++--
 2 files changed, 72 insertions(+), 29 deletions(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index a3b0b3ab28..f137902f1b 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -21,7 +21,7 @@ import json
 import os
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict, Iterable, List, Literal, Optional
+from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
 
 from fastapi import (
     APIRouter,
@@ -147,12 +147,12 @@ if MCP_AVAILABLE:
         user_api_key_dict: UserAPIKeyAuth,
         team_id: Optional[str] = None,
         server_id: Optional[str] = None,
-    ) -> str:
+    ) -> Tuple[str, "LiteLLM_TeamTableCachedObj"]:
         """
         Verify that the caller is an MCP server manager for a team and (for edit/delete)
         that the target server belongs to that team.
 
-        Returns the team_id the caller is managing.
+        Returns a tuple of (team_id, team_obj) for downstream use.
         Raises HTTPException(400) if no team_id can be determined.
         Raises HTTPException(403) if the caller is not an MCP manager or server not in team.
         """
@@ -165,6 +165,7 @@ if MCP_AVAILABLE:
                 detail={"error": "team_id is required for MCP server manager operations."},
             )
 
+        # When the API key is team-scoped, ensure the request team_id matches
         if (
             team_id
             and user_api_key_dict.team_id
@@ -200,7 +201,7 @@ if MCP_AVAILABLE:
                     },
                 )
 
-        return resolved_team_id
+        return resolved_team_id, team_obj
 
     @dataclass
     class _TemporaryMCPServerEntry:
@@ -1277,6 +1278,7 @@ if MCP_AVAILABLE:
         is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role
         manager_team_id: Optional[str] = None
 
+        manager_team_obj = None
         if not is_proxy_admin:
             # Check if the user is an MCP manager for the specified team
             if payload.team_id is None:
@@ -1286,7 +1288,7 @@ if MCP_AVAILABLE:
                         "error": "team_id is required when creating MCP servers as a team MCP manager."
                     },
                 )
-            manager_team_id = await _assert_can_manage_team_mcp_server(
+            manager_team_id, manager_team_obj = await _assert_can_manage_team_mcp_server(
                 user_api_key_dict=user_api_key_dict,
                 team_id=payload.team_id,
             )
@@ -1334,25 +1336,26 @@ if MCP_AVAILABLE:
             # Ensure registry is up to date by reloading from database
             await global_mcp_server_manager.reload_servers_from_database()
 
-            # If created by an MCP manager, auto-assign the server to their team
-            if manager_team_id is not None:
-                from litellm.proxy.proxy_server import user_api_key_cache
+        except Exception as e:
+            verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}")
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail={"error": f"Error creating mcp server: {str(e)}"},
+            )
 
-                team_obj = await get_team_object(
-                    team_id=manager_team_id,
-                    prisma_client=prisma_client,
-                    user_api_key_cache=user_api_key_cache,
-                    check_db_only=False,
-                )
+        # Auto-assign the server to the manager's team (separate from create
+        # so a failure here doesn't mask the successfully created server).
+        if manager_team_id is not None and manager_team_obj is not None:
+            try:
                 existing_permission_id = getattr(
-                    team_obj, "object_permission_id", None
+                    manager_team_obj, "object_permission_id", None
                 )
 
                 # Read existing mcp_servers list and append the new server
                 existing_mcp_servers: list = []
-                if team_obj.object_permission is not None:
+                if manager_team_obj.object_permission is not None:
                     existing_mcp_servers = (
-                        team_obj.object_permission.mcp_servers or []
+                        manager_team_obj.object_permission.mcp_servers or []
                     )
                 updated_mcp_servers = list(
                     set(existing_mcp_servers + [new_mcp_server.server_id])
@@ -1374,12 +1377,10 @@ if MCP_AVAILABLE:
                         where={"team_id": manager_team_id},
                         data={"object_permission_id": new_permission_id},
                     )
-        except Exception as e:
-            verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}")
-            raise HTTPException(
-                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail={"error": f"Error creating mcp server: {str(e)}"},
-            )
+            except Exception as e:
+                verbose_proxy_logger.exception(
+                    f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}"
+                )
         return _redact_mcp_credentials(new_mcp_server)
 
     @router.post(
@@ -1593,8 +1594,9 @@ if MCP_AVAILABLE:
         )
 
         # Authz - proxy admins or team MCP managers can delete MCP servers
+        manager_team_obj = None
         if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
-            await _assert_can_manage_team_mcp_server(
+            _, manager_team_obj = await _assert_can_manage_team_mcp_server(
                 user_api_key_dict=user_api_key_dict,
                 server_id=server_id,
             )
@@ -1612,6 +1614,36 @@ if MCP_AVAILABLE:
         # Ensure registry is up to date by reloading from database
         await global_mcp_server_manager.reload_servers_from_database()
 
+        # Remove server from the manager's team permission list
+        if manager_team_obj is not None:
+            try:
+                existing_permission_id = getattr(
+                    manager_team_obj, "object_permission_id", None
+                )
+                if (
+                    existing_permission_id is not None
+                    and manager_team_obj.object_permission is not None
+                ):
+                    existing_mcp_servers = (
+                        manager_team_obj.object_permission.mcp_servers or []
+                    )
+                    updated_mcp_servers = [
+                        s for s in existing_mcp_servers if s != server_id
+                    ]
+                    await handle_update_object_permission_common(
+                        data_json={
+                            "object_permission": {
+                                "mcp_servers": updated_mcp_servers,
+                            }
+                        },
+                        existing_object_permission_id=existing_permission_id,
+                        prisma_client=prisma_client,
+                    )
+            except Exception as e:
+                verbose_proxy_logger.exception(
+                    f"MCP server deleted but failed to remove from team permissions: {str(e)}"
+                )
+
         # TODO: Enterprise: Finish audit log trail
         if litellm.store_audit_logs:
             pass
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
index 860b255f9c..f4e91157f2 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -9,6 +9,9 @@ from litellm.proxy._types import (
 from litellm.proxy.management_endpoints.common_utils import (
     _is_user_team_mcp_manager,
 )
+from litellm.proxy.management_endpoints import (
+    mcp_management_endpoints as mgmt_endpoints,
+)
 
 
 class TestIsUserTeamMcpManager:
@@ -70,6 +73,9 @@ from litellm.proxy._types import LiteLLM_TeamTableCachedObj
 
 
 @pytest.mark.asyncio
+@pytest.mark.skipif(
+    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
+)
 class TestAssertCanManageTeamMcpServer:
     async def test_mcp_manager_with_team_id_succeeds(self):
         from litellm.proxy.management_endpoints.mcp_management_endpoints import (
@@ -86,10 +92,11 @@ class TestAssertCanManageTeamMcpServer:
             "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
             AsyncMock(return_value=mock_team),
         ):
-            result = await _assert_can_manage_team_mcp_server(
+            team_id, team_obj = await _assert_can_manage_team_mcp_server(
                 user_api_key_dict=user_auth, team_id="team1"
             )
-            assert result == "team1"
+            assert team_id == "team1"
+            assert team_obj == mock_team
 
     async def test_regular_user_gets_403(self):
         from litellm.proxy.management_endpoints.mcp_management_endpoints import (
@@ -154,10 +161,11 @@ class TestAssertCanManageTeamMcpServer:
                 AsyncMock(return_value={"server1", "server2"}),
             ),
         ):
-            result = await _assert_can_manage_team_mcp_server(
+            team_id, team_obj = await _assert_can_manage_team_mcp_server(
                 user_api_key_dict=user_auth, server_id="server1"
             )
-            assert result == "team1"
+            assert team_id == "team1"
+            assert team_obj == mock_team
 
     async def test_mcp_manager_server_not_in_team_gets_403(self):
         from litellm.proxy.management_endpoints.mcp_management_endpoints import (
@@ -188,6 +196,9 @@ class TestAssertCanManageTeamMcpServer:
 
 
 @pytest.mark.asyncio
+@pytest.mark.skipif(
+    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
+)
 class TestCreateMcpServerAsManager:
     async def test_create_auto_assigns_to_team(self):
         """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable."""
@@ -234,7 +245,7 @@ class TestCreateMcpServerAsManager:
             ),
             patch(
                 "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(return_value="team1"),
+                AsyncMock(return_value=("team1", mock_team)),
             ),
             patch(
                 "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",

From 65bfd449e980ac08e6339af7a569a69dfa38b192 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 23:23:18 -0700
Subject: [PATCH 20/31] address greptile review feedback (greploop iteration 2)

- Remove dead get_team_object mock in test (now reuses team_obj from _assert)
- Add test for existing_permission_id=None branch (team linkage)
- Remaining P1s are by-design per spec (admin blocked from MCP, team_id
  ignored for proxy admin)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../test_mcp_manager_role.py                  | 81 ++++++++++++++++++-
 1 file changed, 77 insertions(+), 4 deletions(-)

diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
index f4e91157f2..67e9dda4ce 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -262,10 +262,6 @@ class TestCreateMcpServerAsManager:
                     reload_servers_from_database=AsyncMock(),
                 ),
             ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
-                AsyncMock(return_value=mock_team),
-            ),
             patch(
                 "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
                 mock_handle_update,
@@ -280,3 +276,80 @@ class TestCreateMcpServerAsManager:
             assert "existing_server" in mcp_servers
             assert "new_server_id" in mcp_servers
             assert call_kwargs["existing_object_permission_id"] == "perm1"
+
+    async def test_create_links_new_permission_to_team_when_none_exists(self):
+        """When team has no object_permission_id, create should link the new one."""
+        from litellm.proxy._types import NewMCPServerRequest
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            add_mcp_server,
+        )
+
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+            team_id="team1",
+        )
+
+        payload = NewMCPServerRequest(
+            server_name="test_server",
+            url="https://example.com/mcp",
+            team_id="team1",
+        )
+
+        # Team with NO object_permission_id
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager"),
+            ],
+            object_permission_id=None,
+        )
+        mock_team.object_permission = None
+
+        created_server = MagicMock()
+        created_server.server_id = "new_server_id"
+        created_server.credentials = None
+
+        mock_prisma = MagicMock()
+        mock_prisma.db.litellm_teamtable.update = AsyncMock()
+
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
+                return_value=mock_prisma,
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+                AsyncMock(return_value=("team1", mock_team)),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
+                AsyncMock(return_value=None),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server",
+                AsyncMock(return_value=created_server),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+                MagicMock(
+                    add_server=AsyncMock(),
+                    reload_servers_from_database=AsyncMock(),
+                ),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
+                AsyncMock(return_value="new_perm_id"),
+            ),
+        ):
+            await add_mcp_server(payload=payload, user_api_key_dict=user_auth)
+
+            # Verify the team was updated with the new object_permission_id
+            mock_prisma.db.litellm_teamtable.update.assert_called_once_with(
+                where={"team_id": "team1"},
+                data={"object_permission_id": "new_perm_id"},
+            )

From 34d954f8cb9d02c5c4176368a9baf0dc0dbb5e8d Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Thu, 19 Mar 2026 23:29:38 -0700
Subject: [PATCH 21/31] [Fix] UI: Migrate AntD message API to use context-based
 MessageManager

AntD v5 static message API doesn't render without an App wrapper or
useMessage() context holder. Mirrors the existing notification pattern
by adding message.useMessage() to AntdGlobalProvider and routing all
calls through a new MessageManager module.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../AccessGroupCreateModal.tsx                |  5 ++-
 .../AccessGroupEditModal.tsx                  |  5 ++-
 .../src/components/BulkEditUsers.tsx          |  4 +-
 .../CloudZeroCreateModal.tsx                  |  9 +++--
 .../CloudZeroIntegrationSettings.tsx          | 15 ++++----
 .../CloudZeroUpdateModal.tsx                  |  9 +++--
 .../ProjectModals/CreateProjectModal.tsx      |  7 ++--
 .../ProjectModals/EditProjectModal.tsx        |  7 ++--
 .../SearchTools/SearchToolTester.tsx          |  5 ++-
 .../RouterSettings/Fallbacks/AddFallbacks.tsx |  5 ++-
 .../Fallbacks/FallbackSelectionForm.tsx       |  5 ++-
 .../src/components/agents/add_agent_form.tsx  | 11 +++---
 .../src/components/agents/agent_info.tsx      |  9 +++--
 .../src/components/chat/ChatPage.tsx          |  5 ++-
 .../src/components/chat/MCPAppsPanel.tsx      |  6 +--
 .../src/components/chat/MCPConnectPicker.tsx  |  7 ++--
 .../src/components/chat/MCPCredentialsTab.tsx |  5 ++-
 .../claude_code_plugins/add_plugin_form.tsx   | 17 +++++----
 .../mcp_tools/ByokCredentialModal.tsx         |  9 +++--
 .../components/molecules/message_manager.tsx  | 38 +++++++++++++++++++
 .../src/components/networking.tsx             |  4 +-
 .../organisms/create_key_button.tsx           |  2 +-
 .../chat_ui/CodeInterpreterTool.tsx           |  5 ++-
 .../src/components/policies/index.tsx         | 29 +++++++-------
 .../policies/pipeline_flow_builder.tsx        |  9 +++--
 .../components/policies/policy_templates.tsx  |  5 ++-
 .../shared/CreatedKeyDisplay.test.tsx         | 14 +++----
 .../components/shared/CreatedKeyDisplay.tsx   |  5 ++-
 .../src/components/team/TeamInfo.tsx          |  7 ++--
 .../CreateVectorStore.tsx                     | 19 +++++-----
 .../DocumentsTable.tsx                        |  5 ++-
 .../VectorStoreTester.tsx                     |  5 ++-
 .../view_logs/LogDetailsDrawer/InputCard.tsx  |  4 +-
 .../view_logs/LogDetailsDrawer/OutputCard.tsx |  5 ++-
 .../src/contexts/AntdGlobalProvider.tsx       | 14 ++++---
 35 files changed, 190 insertions(+), 125 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/molecules/message_manager.tsx

diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx
index b51fbd5bd9..5ec9f04f3b 100644
--- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx
+++ b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx
@@ -1,5 +1,6 @@
 import React from "react";
-import { Modal, Form, message } from "antd";
+import { Modal, Form } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   AccessGroupBaseForm,
   AccessGroupFormValues,
@@ -37,7 +38,7 @@ export function AccessGroupCreateModal({
 
         createMutation.mutate(params, {
           onSuccess: () => {
-            message.success("Access group created successfully");
+            MessageManager.success("Access group created successfully");
             form.resetFields();
             onSuccess?.();
             onCancel();
diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx
index 919295b6f7..f05edb3fe9 100644
--- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx
+++ b/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx
@@ -1,5 +1,6 @@
 import React, { useEffect } from "react";
-import { Modal, Form, message } from "antd";
+import { Modal, Form } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   AccessGroupBaseForm,
   AccessGroupFormValues,
@@ -55,7 +56,7 @@ export function AccessGroupEditModal({
           { accessGroupId: accessGroup.access_group_id, params },
           {
             onSuccess: () => {
-              message.success("Access group updated successfully");
+              MessageManager.success("Access group updated successfully");
               onSuccess?.();
               onCancel();
             },
diff --git a/ui/litellm-dashboard/src/components/BulkEditUsers.tsx b/ui/litellm-dashboard/src/components/BulkEditUsers.tsx
index 2f3e57ff2a..2ca69065bc 100644
--- a/ui/litellm-dashboard/src/components/BulkEditUsers.tsx
+++ b/ui/litellm-dashboard/src/components/BulkEditUsers.tsx
@@ -3,7 +3,6 @@ import {
   Modal,
   Typography,
   Divider,
-  message,
   Table,
   Select,
   InputNumber,
@@ -14,6 +13,7 @@ import {
 import { userBulkUpdateUserCall, teamBulkMemberAddCall, Member } from "./networking";
 import { UserEditView } from "./user_edit_view";
 import NotificationsManager from "./molecules/notifications_manager";
+import MessageManager from "@/components/molecules/message_manager";
 
 const { Text, Title } = Typography;
 
@@ -188,7 +188,7 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
         }
 
         if (failedTeams.length > 0) {
-          message.warning(`Failed to add users to ${failedTeams.length} team(s)`);
+          MessageManager.warning(`Failed to add users to ${failedTeams.length} team(s)`);
         }
       }
 
diff --git a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx
index feb00fc040..5f460cb7cb 100644
--- a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx
+++ b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx
@@ -1,4 +1,5 @@
-import { Form, Modal, Input, message } from "antd";
+import { Form, Modal, Input } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useEffect } from "react";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useCloudZeroCreate } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate";
@@ -31,7 +32,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
         },
         {
           onSuccess: () => {
-            message.success("CloudZero integration created successfully");
+            MessageManager.success("CloudZero integration created successfully");
             form.resetFields();
             onOk();
           },
@@ -39,7 +40,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
             if (error?.errorFields) {
               return;
             }
-            message.error(error?.message || "Failed to create CloudZero integration");
+            MessageManager.error(error?.message || "Failed to create CloudZero integration");
           },
         },
       );
@@ -47,7 +48,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
       if (error?.errorFields) {
         return;
       }
-      message.error(error?.message || "Failed to create CloudZero integration");
+      MessageManager.error(error?.message || "Failed to create CloudZero integration");
     }
   };
 
diff --git a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx
index c161d241f7..62709dac52 100644
--- a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx
+++ b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx
@@ -3,7 +3,8 @@ import { useCloudZeroExport } from "@/app/(dashboard)/hooks/cloudzero/useCloudZe
 import { useCloudZeroDeleteSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
-import { Alert, Button, Card, Descriptions, Divider, message, Popconfirm, Tag } from "antd";
+import { Alert, Button, Card, Descriptions, Divider, Popconfirm, Tag } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CheckCircle, Edit, Play, Trash2, Upload } from "lucide-react";
 import { useState } from "react";
 import CloudZeroUpdateModal from "./CloudZeroUpdateModal";
@@ -30,10 +31,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
       { limit: 10 },
       {
         onSuccess: (data) => {
-          message.success("Dry run completed successfully");
+          MessageManager.success("Dry run completed successfully");
         },
         onError: (error) => {
-          message.error(error?.message || "Failed to perform dry run");
+          MessageManager.error(error?.message || "Failed to perform dry run");
         },
       },
     );
@@ -48,10 +49,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
       { operation: "replace_hourly" },
       {
         onSuccess: () => {
-          message.success("Data successfully exported to CloudZero");
+          MessageManager.success("Data successfully exported to CloudZero");
         },
         onError: (error) => {
-          message.error(error?.message || "Failed to export data");
+          MessageManager.error(error?.message || "Failed to export data");
         },
       },
     );
@@ -79,12 +80,12 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
 
     deleteMutation.mutate(undefined, {
       onSuccess: () => {
-        message.success("CloudZero integration deleted successfully");
+        MessageManager.success("CloudZero integration deleted successfully");
         setIsDeleteModalOpen(false);
         onSettingsUpdated();
       },
       onError: (error) => {
-        message.error(error?.message || "Failed to delete CloudZero integration");
+        MessageManager.error(error?.message || "Failed to delete CloudZero integration");
       },
     });
   };
diff --git a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx
index 0aca6857b8..a04007897c 100644
--- a/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx
+++ b/ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx
@@ -1,6 +1,7 @@
 import { useCloudZeroUpdateSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import { Form, Input, message, Modal } from "antd";
+import { Form, Input, Modal } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useEffect } from "react";
 import { CloudZeroSettings } from "./types";
 
@@ -39,7 +40,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
         },
         {
           onSuccess: () => {
-            message.success("CloudZero integration updated successfully");
+            MessageManager.success("CloudZero integration updated successfully");
             form.resetFields();
             onOk();
           },
@@ -47,7 +48,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
             if (error?.errorFields) {
               return;
             }
-            message.error(error?.message || "Failed to update CloudZero integration");
+            MessageManager.error(error?.message || "Failed to update CloudZero integration");
           },
         },
       );
@@ -55,7 +56,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
       if (error?.errorFields) {
         return;
       }
-      message.error(error?.message || "Failed to update CloudZero integration");
+      MessageManager.error(error?.message || "Failed to update CloudZero integration");
     }
   };
 
diff --git a/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx b/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx
index e490f89303..bbf56e4930 100644
--- a/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx
+++ b/ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx
@@ -1,5 +1,6 @@
-import { Modal, Form, Button, Typography, message } from "antd";
+import { Modal, Form, Button, Typography } from "antd";
 import { FolderAddOutlined } from "@ant-design/icons";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   useCreateProject,
   ProjectCreateParams,
@@ -32,12 +33,12 @@ export function CreateProjectModal({
 
       createMutation.mutate(params, {
         onSuccess: () => {
-          message.success("Project created successfully");
+          MessageManager.success("Project created successfully");
           form.resetFields();
           onClose();
         },
         onError: (error) => {
-          message.error(error.message || "Failed to create project");
+          MessageManager.error(error.message || "Failed to create project");
         },
       });
     } catch (error) {
diff --git a/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx
index 75f56b1373..dc3b43ef73 100644
--- a/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx
+++ b/ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx
@@ -1,6 +1,7 @@
 import { useEffect } from "react";
-import { Modal, Form, Button, Typography, message } from "antd";
+import { Modal, Form, Button, Typography } from "antd";
 import { SaveOutlined } from "@ant-design/icons";
+import MessageManager from "@/components/molecules/message_manager";
 import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects";
 import {
   useUpdateProject,
@@ -80,12 +81,12 @@ export function EditProjectModal({
         { projectId: project.project_id, params },
         {
           onSuccess: () => {
-            message.success("Project updated successfully");
+            MessageManager.success("Project updated successfully");
             onSuccess?.();
             onClose();
           },
           onError: (error) => {
-            message.error(error.message || "Failed to update project");
+            MessageManager.error(error.message || "Failed to update project");
           },
         },
       );
diff --git a/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx b/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx
index 28b34b0608..d1cb5077b1 100644
--- a/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx
+++ b/ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx
@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Button, Input, Typography, Spin, message } from "antd";
+import { Button, Input, Typography, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { SearchOutlined, LoadingOutlined } from "@ant-design/icons";
 import { searchToolQueryCall } from "../networking";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -39,7 +40,7 @@ export const SearchToolTester: React.FC<SearchToolTesterProps> = ({ searchToolNa
 
   const handleSearch = async () => {
     if (!query.trim()) {
-      message.warning("Please enter a search query");
+      MessageManager.warning("Please enter a search query");
       return;
     }
 
diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx
index 34f059516a..e55089d27d 100644
--- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx
@@ -5,8 +5,9 @@
  */
 
 import { Button as TremorButton } from "@tremor/react";
-import { Button, message } from "antd";
+import { Button } from "antd";
 import React, { useEffect, useState } from "react";
+import MessageManager from "@/components/molecules/message_manager";
 import NotificationManager from "../../../molecules/notifications_manager";
 import { fetchAvailableModels, ModelGroup } from "../../../playground/llm_calls/fetch_models";
 import { AddFallbacksModal } from "./AddFallbacksModal";
@@ -90,7 +91,7 @@ export default function AddFallbacks({
       (g) => !g.primaryModel || g.fallbackModels.length === 0,
     );
     if (invalidGroups.length > 0) {
-      message.error(
+      MessageManager.error(
         `Please complete configuration for all groups. ${invalidGroups.length} group(s) incomplete.`,
       );
       return;
diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx
index 08b031c683..0482161bc6 100644
--- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx
@@ -5,9 +5,10 @@
  */
 
 import { Button } from "@tremor/react";
-import { message, Tabs } from "antd";
+import { Tabs } from "antd";
 import { Plus } from "lucide-react";
 import React, { useEffect, useState } from "react";
+import MessageManager from "@/components/molecules/message_manager";
 import { FallbackGroup, FallbackGroupConfig } from "./FallbackGroupConfig";
 
 interface FallbackSelectionFormProps {
@@ -60,7 +61,7 @@ export function FallbackSelectionForm({
 
   const handleRemoveGroup = (targetId: string) => {
     if (groups.length === 1) {
-      message.warning("At least one group is required");
+      MessageManager.warning("At least one group is required");
       return;
     }
     const newGroups = groups.filter((g) => g.id !== targetId);
diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx
index c5518596b8..b6d96f0445 100644
--- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx
+++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx
@@ -1,5 +1,6 @@
 import React, { useState, useEffect } from "react";
-import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd";
+import { Modal, Form, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button } from "@tremor/react";
 import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
@@ -216,7 +217,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
 
   const handleCreateAgent = async () => {
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
@@ -226,7 +227,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
       const values = { ...form.getFieldsValue(true) };
       const agentData = buildAgentData(values);
       if (!agentData) {
-        message.error("Failed to build agent data");
+        MessageManager.error("Failed to build agent data");
         setIsSubmitting(false);
         return;
       }
@@ -301,7 +302,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
         setCreatedKeyValue(keyResponse.key || null);
       } else if (keyAssignOption === "existing_key") {
         if (!selectedExistingKey) {
-          message.error("Please select an existing key to assign");
+          MessageManager.error("Please select an existing key to assign");
           setIsSubmitting(false);
           return;
         }
@@ -318,7 +319,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
     } catch (error) {
       console.error("Error creating agent:", error);
       const errorMessage = error instanceof Error ? error.message : String(error);
-      message.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent");
+      MessageManager.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent");
     } finally {
       setIsSubmitting(false);
     }
diff --git a/ui/litellm-dashboard/src/components/agents/agent_info.tsx b/ui/litellm-dashboard/src/components/agents/agent_info.tsx
index b41e318a76..d543be8356 100644
--- a/ui/litellm-dashboard/src/components/agents/agent_info.tsx
+++ b/ui/litellm-dashboard/src/components/agents/agent_info.tsx
@@ -1,6 +1,7 @@
 import React, { useState, useEffect } from "react";
 import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react";
-import { Form, Input, InputNumber, Button as AntButton, message, Spin, Descriptions, Divider } from "antd";
+import { Form, Input, InputNumber, Button as AntButton, Spin, Descriptions, Divider } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking";
 import { Agent } from "./types";
@@ -72,7 +73,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
       }
     } catch (error) {
       console.error("Error fetching agent info:", error);
-      message.error("Failed to load agent information");
+      MessageManager.error("Failed to load agent information");
     } finally {
       setIsLoading(false);
     }
@@ -111,12 +112,12 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
       }
       
       await patchAgentCall(accessToken, agentId, updateData);
-      message.success("Agent updated successfully");
+      MessageManager.success("Agent updated successfully");
       setIsEditing(false);
       fetchAgentInfo();
     } catch (error) {
       console.error("Error updating agent:", error);
-      message.error("Failed to update agent");
+      MessageManager.error("Failed to update agent");
     } finally {
       setIsSaving(false);
     }
diff --git a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx
index ccf39d2147..b547f74917 100644
--- a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx
+++ b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx
@@ -1,7 +1,8 @@
 "use client";
 
 import React, { useCallback, useEffect, useRef, useState, useLayoutEffect } from "react";
-import { Tooltip, Skeleton, Popover, message } from "antd";
+import { Tooltip, Skeleton, Popover } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   SettingOutlined,
   PlusOutlined,
@@ -212,7 +213,7 @@ const ChatPage: React.FC<ChatPageProps> = ({ accessToken, userRole, userId, user
           localStorage.setItem(LOCALSTORAGE_MODEL_KEY, JSON.stringify([names[0]]));
         }
       })
-      .catch(() => message.error("Could not load models"))
+      .catch(() => MessageManager.error("Could not load models"))
       .finally(() => setIsLoadingModels(false));
   }, [accessToken]);
 
diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
index b805cab71e..d25db73ae4 100644
--- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
+++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx
@@ -5,7 +5,7 @@ import { Spin, Input, Button, Skeleton } from "antd";
 import { SearchOutlined, ArrowLeftOutlined, RightOutlined, ToolOutlined, CheckCircleOutlined } from "@ant-design/icons";
 import { deleteMCPOAuthUserCredential, fetchMCPServers, getMCPOAuthUserCredentialStatus, listMCPTools } from "../networking";
 import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types";
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
 
 // ── OAuth2 connect button ─────────────────────────────────────────────────────
@@ -198,7 +198,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
       const idToFetch = serverId ?? serverName;
       const result = await listMCPTools(accessToken, idToFetch);
       if (result?.error) {
-        message.warning(`Could not load tools for ${serverName}`);
+        MessageManager.warning(`Could not load tools for ${serverName}`);
         return;
       }
       // Use the ref so we read the most up-to-date list; guard against duplicates
@@ -207,7 +207,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
         onChange([...selectedServersRef.current, serverName]);
       }
     } catch {
-      message.warning(`Could not load tools for ${serverName}`);
+      MessageManager.warning(`Could not load tools for ${serverName}`);
     } finally {
       setTogglingOn((prev) => {
         const next = new Set(prev);
diff --git a/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx b/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx
index a52ce18156..6ef3aecc46 100644
--- a/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx
+++ b/ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx
@@ -1,5 +1,6 @@
 import React, { useEffect, useState } from "react";
-import { Switch, Spin, message } from "antd";
+import { Switch, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { fetchMCPServers, listMCPTools } from "../networking";
 import { MCPServer } from "../mcp_tools/types";
 
@@ -57,7 +58,7 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
       const result = await listMCPTools(accessToken, serverName);
       // listMCPTools never throws; it returns { tools, error, message } on failure
       if (result?.error) {
-        message.warning(
+        MessageManager.warning(
           `Could not load tools for ${serverName} — it will be excluded from this message.`
         );
         // Do not add to selectedServers
@@ -65,7 +66,7 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
       }
       onChange([...selectedServers, serverName]);
     } catch {
-      message.warning(
+      MessageManager.warning(
         `Could not load tools for ${serverName} — it will be excluded from this message.`
       );
       // Do not add to selectedServers
diff --git a/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx b/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx
index 363ec8c0e4..e2d879b22a 100644
--- a/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx
+++ b/ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx
@@ -8,7 +8,8 @@
  */
 
 import React, { useCallback, useEffect, useState } from "react";
-import { Spin, message } from "antd";
+import { Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { DeleteOutlined, LinkOutlined } from "@ant-design/icons";
 import { Badge, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react";
 import {
@@ -77,7 +78,7 @@ const MCPCredentialsTab: React.FC<Props> = ({ accessToken }) => {
       await deleteMCPOAuthUserCredential(accessToken, serverId);
       setCredentials((prev) => prev.filter((c) => c.server_id !== serverId));
     } catch {
-      message.error("Failed to revoke connection. Please try again.");
+      MessageManager.error("Failed to revoke connection. Please try again.");
     } finally {
       setRevoking((prev) => { const n = new Set(prev); n.delete(serverId); return n; });
     }
diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx
index 217851f128..d5e417a9a8 100644
--- a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx
+++ b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx
@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Modal, Form, Input, Select, message } from "antd";
+import { Modal, Form, Input, Select } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button } from "@tremor/react";
 import { registerClaudeCodePlugin } from "../networking";
 import {
@@ -43,13 +44,13 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
   const handleSubmit = async (values: any) => {
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
     // Validate plugin name
     if (!validatePluginName(values.name)) {
-      message.error(
+      MessageManager.error(
         "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)"
       );
       return;
@@ -57,7 +58,7 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
     // Validate semantic version if provided
     if (values.version && !isValidSemanticVersion(values.version)) {
-      message.error(
+      MessageManager.error(
         "Version must be in semantic versioning format (e.g., 1.0.0)"
       );
       return;
@@ -65,13 +66,13 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
     // Validate email if provided
     if (values.authorEmail && !isValidEmail(values.authorEmail)) {
-      message.error("Invalid email format");
+      MessageManager.error("Invalid email format");
       return;
     }
 
     // Validate homepage URL if provided
     if (values.homepage && !isValidUrl(values.homepage)) {
-      message.error("Invalid homepage URL format");
+      MessageManager.error("Invalid homepage URL format");
       return;
     }
 
@@ -119,14 +120,14 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
       }
 
       await registerClaudeCodePlugin(accessToken, pluginData);
-      message.success("Plugin registered successfully");
+      MessageManager.success("Plugin registered successfully");
       form.resetFields();
       setSourceType("github");
       onSuccess();
       onClose();
     } catch (error) {
       console.error("Error registering plugin:", error);
-      message.error("Failed to register plugin");
+      MessageManager.error("Failed to register plugin");
     } finally {
       setIsSubmitting(false);
     }
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx
index 473918c126..58c2a965e0 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx
@@ -1,7 +1,8 @@
 "use client";
 
 import React, { useState } from "react";
-import { Modal, Input, Switch, message } from "antd";
+import { Modal, Input, Switch } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   KeyOutlined,
   LockOutlined,
@@ -46,7 +47,7 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
 
   const handleAuthorize = async () => {
     if (!apiKey.trim()) {
-      message.error("Please enter your API key");
+      MessageManager.error("Please enter your API key");
       return;
     }
     setLoading(true);
@@ -63,11 +64,11 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
         const err = await response.json();
         throw new Error(err?.detail?.error || "Failed to save credential");
       }
-      message.success(`Connected to ${serverDisplayName}`);
+      MessageManager.success(`Connected to ${serverDisplayName}`);
       onSuccess(server.server_id);
       handleClose();
     } catch (e: any) {
-      message.error(e.message || "Failed to connect");
+      MessageManager.error(e.message || "Failed to connect");
     } finally {
       setLoading(false);
     }
diff --git a/ui/litellm-dashboard/src/components/molecules/message_manager.tsx b/ui/litellm-dashboard/src/components/molecules/message_manager.tsx
new file mode 100644
index 0000000000..e4c1552d5e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/molecules/message_manager.tsx
@@ -0,0 +1,38 @@
+import { message as staticMessage } from "antd";
+import type { MessageInstance } from "antd/es/message/interface";
+
+let messageInstance: MessageInstance | null = null;
+
+export const setMessageInstance = (instance: MessageInstance) => {
+  messageInstance = instance;
+};
+
+const getMessageApi = () => messageInstance || staticMessage;
+
+const MessageManager = {
+  success(content: string, duration?: number) {
+    getMessageApi().success(content, duration);
+  },
+
+  error(content: string, duration?: number) {
+    getMessageApi().error(content, duration);
+  },
+
+  warning(content: string, duration?: number) {
+    getMessageApi().warning(content, duration);
+  },
+
+  info(content: string, duration?: number) {
+    getMessageApi().info(content, duration);
+  },
+
+  loading(content: string, duration?: number) {
+    return getMessageApi().loading(content, duration);
+  },
+
+  destroy() {
+    getMessageApi().destroy();
+  },
+};
+
+export default MessageManager;
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index bb9a0c1e01..956968dfee 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -68,7 +68,7 @@ export const getInProductNudgesCall = async (accessToken: string) => {
 /**
  * Helper file for calls being made to proxy
  */
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types";
 import { Team } from "./key_team_helpers/key_list";
@@ -555,7 +555,7 @@ export const modelCreateCall = async (accessToken: string, formValues: Model) =>
     console.log("API Response:", data);
 
     // Close any existing messages before showing new ones
-    message.destroy();
+    MessageManager.destroy();
 
     // Sequential success messages
     NotificationsManager.success(`Model ${formValues.model_name} created successfully`);
diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
index d0a7a909d6..df5b3a4b32 100644
--- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
+++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
@@ -8,7 +8,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { useQueryClient } from "@tanstack/react-query";
 import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Button as Button2, Form, Input, message, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
+import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
 import debounce from "lodash/debounce";
 import React, { useCallback, useEffect, useState } from "react";
 import { rolesWithWriteAccess } from "../../utils/roles";
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx
index 04711ede9f..dd37769556 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx
@@ -1,5 +1,6 @@
 import React from "react";
-import { Switch, Tooltip, message } from "antd";
+import { Switch, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CodeOutlined, InfoCircleOutlined, ExclamationCircleOutlined } from "@ant-design/icons";
 import { Text } from "@tremor/react";
 
@@ -38,7 +39,7 @@ const CodeInterpreterTool: React.FC<CodeInterpreterToolProps> = ({
 
   const handleToggle = (checked: boolean) => {
     if (checked && !isOpenAI) {
-      message.warning("Code Interpreter is only available for OpenAI models");
+      MessageManager.warning("Code Interpreter is only available for OpenAI models");
       return;
     }
     onEnabledChange(checked);
diff --git a/ui/litellm-dashboard/src/components/policies/index.tsx b/ui/litellm-dashboard/src/components/policies/index.tsx
index bd77f3ab9b..f47b8d78b9 100644
--- a/ui/litellm-dashboard/src/components/policies/index.tsx
+++ b/ui/litellm-dashboard/src/components/policies/index.tsx
@@ -1,6 +1,7 @@
 import React, { useState, useEffect, useCallback } from "react";
 import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
-import { Modal, message, Alert } from "antd";
+import { Modal, Alert } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { ExclamationCircleOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import { isAdminRole } from "@/utils/roles";
 import PolicyTable from "./policy_table";
@@ -80,7 +81,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setPoliciesList(response.policies || []);
     } catch (error) {
       console.error("Error fetching policies:", error);
-      message.error("Failed to fetch policies");
+      MessageManager.error("Failed to fetch policies");
     } finally {
       setIsLoading(false);
     }
@@ -95,7 +96,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setAttachmentsList(response.attachments || []);
     } catch (error) {
       console.error("Error fetching attachments:", error);
-      message.error("Failed to fetch attachments");
+      MessageManager.error("Failed to fetch attachments");
     } finally {
       setIsAttachmentsLoading(false);
     }
@@ -148,11 +149,11 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
     setIsDeleting(true);
     try {
       await deletePolicyCall(accessToken, policyToDelete.policy_id);
-      message.success(`Policy "${policyToDelete.policy_name}" deleted successfully`);
+      MessageManager.success(`Policy "${policyToDelete.policy_name}" deleted successfully`);
       await fetchPolicies();
     } catch (error) {
       console.error("Error deleting policy:", error);
-      message.error("Failed to delete policy");
+      MessageManager.error("Failed to delete policy");
     } finally {
       setIsDeleting(false);
       setIsDeleteModalOpen(false);
@@ -177,11 +178,11 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
         if (!accessToken) return;
         try {
           await deletePolicyAttachmentCall(accessToken, attachmentId);
-          message.success("Attachment deleted successfully");
+          MessageManager.success("Attachment deleted successfully");
           fetchAttachments();
         } catch (error) {
           console.error("Error deleting attachment:", error);
-          message.error("Failed to delete attachment");
+          MessageManager.error("Failed to delete attachment");
         }
       },
     });
@@ -193,7 +194,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
 
   const handleUseTemplate = async (template: any) => {
     if (!accessToken) {
-      message.error("Authentication required");
+      MessageManager.error("Authentication required");
       return;
     }
 
@@ -221,7 +222,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setIsGuardrailSelectionModalOpen(true);
     } catch (error) {
       console.error("Error fetching guardrails:", error);
-      message.error("Failed to load guardrails. Please try again.");
+      MessageManager.error("Failed to load guardrails. Please try again.");
     }
   };
 
@@ -271,7 +272,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       await proceedWithTemplate(enrichedTemplate);
     } catch (error) {
       console.error("Error enriching template:", error);
-      message.error("Failed to configure template. Please try again.");
+      MessageManager.error("Failed to configure template. Please try again.");
       setIsEnrichingTemplate(false);
     }
   };
@@ -318,15 +319,15 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
 
       // Show success message
       if (createdGuardrails.length > 0) {
-        message.success(
+        MessageManager.success(
           `Created ${createdGuardrails.length} guardrail${createdGuardrails.length > 1 ? "s" : ""}! Complete the policy form to save.`
         );
       } else {
-        message.success("Template ready! Complete the policy form to save.");
+        MessageManager.success("Template ready! Complete the policy form to save.");
       }
 
       if (failedGuardrails.length > 0) {
-        message.warning(
+        MessageManager.warning(
           `Failed to create ${failedGuardrails.length} guardrail(s): ${failedGuardrails.join(", ")}. You may need to create them manually.`
         );
       }
@@ -348,7 +349,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setTemplateQueue([]);
       setTemplateQueueProgress(null);
       console.error("Error creating guardrails:", error);
-      message.error("Failed to create guardrails. Please try again.");
+      MessageManager.error("Failed to create guardrails. Please try again.");
     }
   };
 
diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx
index 89665774ee..b1768d5b81 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx
@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Select, Typography, message, Spin } from "antd";
+import { Select, Typography, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button, TextInput } from "@tremor/react";
 import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline";
 import { DotsVerticalIcon } from "@heroicons/react/solid";
@@ -1385,17 +1386,17 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
 
   const handleSave = async () => {
     if (!policyName.trim()) {
-      message.error("Please enter a policy name");
+      MessageManager.error("Please enter a policy name");
       return;
     }
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
     const emptySteps = pipeline.steps.filter((s) => !s.guardrail);
     if (emptySteps.length > 0) {
-      message.error("Please select a guardrail for all steps");
+      MessageManager.error("Please select a guardrail for all steps");
       return;
     }
 
diff --git a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx
index 98ba0acb6a..c4ea22651c 100644
--- a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx
+++ b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx
@@ -1,5 +1,6 @@
 import React, { useState, useEffect, useMemo } from "react";
-import { Card, Button, Spin, message, Checkbox } from "antd";
+import { Card, Button, Spin, Checkbox } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   ShieldCheckIcon,
   ShieldExclamationIcon,
@@ -184,7 +185,7 @@ const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, onOpen
         onTemplatesLoaded?.(data);
       } catch (error) {
         console.error("Error fetching policy templates:", error);
-        message.error("Failed to fetch policy templates");
+        MessageManager.error("Failed to fetch policy templates");
       } finally {
         setIsLoading(false);
       }
diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx
index 013e222c6b..eb6fe5f50c 100644
--- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx
+++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx
@@ -3,15 +3,11 @@ import { render, screen, act } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import CreatedKeyDisplay from "./CreatedKeyDisplay";
 
-vi.mock("antd", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("antd")>();
-  return {
-    ...actual,
-    message: { success: vi.fn() },
-  };
-});
+vi.mock("@/components/molecules/message_manager", () => ({
+  default: { success: vi.fn(), error: vi.fn(), warning: vi.fn(), info: vi.fn(), loading: vi.fn(), destroy: vi.fn() },
+}));
 
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 
 describe("CreatedKeyDisplay", () => {
   beforeEach(() => {
@@ -52,7 +48,7 @@ describe("CreatedKeyDisplay", () => {
 
     await user.click(screen.getByRole("button", { name: /copy virtual key/i }));
 
-    expect(message.success).toHaveBeenCalledWith("Key copied to clipboard");
+    expect(MessageManager.success).toHaveBeenCalledWith("Key copied to clipboard");
   });
 
   it("should revert button text back after 2 seconds", async () => {
diff --git a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx
index d75c63c646..00ffb8c6e2 100644
--- a/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx
+++ b/ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx
@@ -1,6 +1,7 @@
 import React, { useState } from "react";
 import { CopyToClipboard } from "react-copy-to-clipboard";
-import { Button, message } from "antd";
+import { Button } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 
 interface CreatedKeyDisplayProps {
   apiKey: string;
@@ -15,7 +16,7 @@ const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
 
   const handleCopy = () => {
     setCopied(true);
-    message.success("Key copied to clipboard");
+    MessageManager.success("Key copied to clipboard");
     setTimeout(() => setCopied(false), 2000);
   };
 
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index d2ce79580d..a4c7ae2bbb 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -20,7 +20,8 @@ import { isProxyAdminRole } from "@/utils/roles";
 import { EditOutlined, InfoCircleOutlined, SaveOutlined } from "@ant-design/icons";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { Badge, Card, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Button, Form, Input, message, Select, Switch, Tabs, Tooltip } from "antd";
+import { Button, Form, Input, Select, Switch, Tabs, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CheckIcon, CopyIcon } from "lucide-react";
 import React, { useEffect, useMemo, useState } from "react";
 import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
@@ -366,7 +367,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
         tpm_limit: values.tpm_limit,
         rpm_limit: values.rpm_limit,
       };
-      message.destroy(); // Remove all existing toasts
+      MessageManager.destroy(); // Remove all existing toasts
 
       await teamMemberUpdateCall(accessToken, teamId, member);
 
@@ -388,7 +389,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
       }
       setIsEditMemberModalVisible(false);
 
-      message.destroy(); // Remove all existing toasts
+      MessageManager.destroy(); // Remove all existing toasts
 
       NotificationsManager.fromBackend(errMsg);
       console.error("Error updating team member:", error);
diff --git a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx
index 97162aa132..b189168e1b 100644
--- a/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx
+++ b/ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx
@@ -1,6 +1,7 @@
 import React, { useState } from "react";
 import { Card, Title, Text } from "@tremor/react";
-import { Upload, Button, Select, Form, message, Alert, Tooltip, Input } from "antd";
+import { Upload, Button, Select, Form, Alert, Tooltip, Input } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { InboxOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import type { UploadProps } from "antd";
 import { ragIngestCall } from "../networking";
@@ -47,13 +48,13 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
       ].includes(file.type);
 
       if (!isValidType) {
-        message.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`);
+        MessageManager.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`);
         return Upload.LIST_IGNORE;
       }
 
       const isLt50M = file.size / 1024 / 1024 < 50;
       if (!isLt50M) {
-        message.error(`${file.name} must be smaller than 50MB!`);
+        MessageManager.error(`${file.name} must be smaller than 50MB!`);
         return Upload.LIST_IGNORE;
       }
 
@@ -87,12 +88,12 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
 
   const handleCreateVectorStore = async () => {
     if (documents.length === 0) {
-      message.warning("Please upload at least one document");
+      MessageManager.warning("Please upload at least one document");
       return;
     }
 
     if (!selectedProvider) {
-      message.warning("Please select a provider");
+      MessageManager.warning("Please select a provider");
       return;
     }
 
@@ -100,7 +101,7 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
     const requiredFields = getProviderSpecificFields(selectedProvider).filter((field) => field.required);
     for (const field of requiredFields) {
       if (!providerParams[field.name]) {
-        message.warning(`Please provide ${field.label}`);
+        MessageManager.warning(`Please provide ${field.label}`);
         return;
       }
     }
@@ -108,17 +109,17 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
     // S3 Vectors specific validation
     if (selectedProvider === "s3_vectors") {
       if (providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3) {
-        message.warning("Vector bucket name must be at least 3 characters");
+        MessageManager.warning("Vector bucket name must be at least 3 characters");
         return;
       }
       if (providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3) {
-        message.warning("Index name must be at least 3 characters if provided");
+        MessageManager.warning("Index name must be at least 3 characters if provided");
         return;
       }
     }
 
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
diff --git a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx
index aeb4240d36..c1ce57ad33 100644
--- a/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx
+++ b/ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx
@@ -1,5 +1,6 @@
 import React from "react";
-import { Table, Badge, Tooltip, message } from "antd";
+import { Table, Badge, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { EyeOutlined, CopyOutlined, DeleteOutlined } from "@ant-design/icons";
 import { DocumentUpload } from "./types";
 
@@ -11,7 +12,7 @@ interface DocumentsTableProps {
 const DocumentsTable: React.FC<DocumentsTableProps> = ({ documents, onRemove }) => {
   const handleCopyId = (uid: string) => {
     navigator.clipboard.writeText(uid);
-    message.success("Document ID copied to clipboard");
+    MessageManager.success("Document ID copied to clipboard");
   };
 
   const getStatusBadge = (status: DocumentUpload["status"]) => {
diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx
index f8b5ada0c3..5c59334fa4 100644
--- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx
+++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx
@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Button, Input, Card, Typography, Spin, message, Divider } from "antd";
+import { Button, Input, Card, Typography, Spin, Divider } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { SendOutlined, DatabaseOutlined, LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
 import { vectorStoreSearchCall } from "../networking";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -46,7 +47,7 @@ export const VectorStoreTester: React.FC<VectorStoreTesterProps> = ({ vectorStor
 
   const handleSearch = async () => {
     if (!query.trim()) {
-      message.warning("Please enter a search query");
+      MessageManager.warning("Please enter a search query");
       return;
     }
 
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx
index 1299f9c168..8fb47274a8 100644
--- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx
@@ -4,7 +4,7 @@
  */
 
 import { useState } from 'react';
-import { message } from 'antd';
+import MessageManager from "@/components/molecules/message_manager";
 import { ParsedMessage } from './prettyMessagesTypes';
 import { SectionHeader } from './SectionHeader';
 import { CollapsibleMessage } from './CollapsibleMessage';
@@ -33,7 +33,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps)
   const handleCopy = () => {
     const content = lastMessage?.content || '';
     navigator.clipboard.writeText(content);
-    message.success('Input copied');
+    MessageManager.success('Input copied');
   };
 
   return (
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx
index eff8d83cbd..22f1708ee4 100644
--- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx
@@ -4,7 +4,8 @@
  */
 
 import { useState } from 'react';
-import { Typography, message as antdMessage } from 'antd';
+import { Typography } from 'antd';
+import MessageManager from "@/components/molecules/message_manager";
 import { ParsedMessage } from './prettyMessagesTypes';
 import { SectionHeader } from './SectionHeader';
 import { SimpleMessageBlock } from './SimpleMessageBlock';
@@ -25,7 +26,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard
     
     const content = message.content || '';
     navigator.clipboard.writeText(content);
-    antdMessage.success('Output copied');
+    MessageManager.success('Output copied');
   };
 
   if (!message) {
diff --git a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx
index 5b5c1036fa..cc182ee36d 100644
--- a/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx
+++ b/ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx
@@ -1,23 +1,27 @@
 "use client";
 
 import React, { useEffect, useRef } from "react";
-import { notification } from "antd";
+import { notification, message } from "antd";
 import { setNotificationInstance } from "@/components/molecules/notifications_manager";
+import { setMessageInstance } from "@/components/molecules/message_manager";
 
 export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) {
-  const [api, contextHolder] = notification.useNotification();
+  const [notificationApi, notificationContextHolder] = notification.useNotification();
+  const [messageApi, messageContextHolder] = message.useMessage();
   const initialized = useRef(false);
 
   useEffect(() => {
     if (!initialized.current) {
-      setNotificationInstance(api);
+      setNotificationInstance(notificationApi);
+      setMessageInstance(messageApi);
       initialized.current = true;
     }
-  }, [api]);
+  }, [notificationApi, messageApi]);
 
   return (
     <>
-      {contextHolder}
+      {notificationContextHolder}
+      {messageContextHolder}
       {children}
     </>
   );

From b8c9bf7d256ed90b1f8cfdbe8b57d79f118009f4 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 09:05:18 -0700
Subject: [PATCH 22/31] refactor: extract _auto_assign and _remove helpers, use
 team_endpoints helper

- Replace raw prisma_client.db.litellm_teamtable.update with
  handle_update_object_permission from team_endpoints (follows
  established helper-function pattern)
- Extract _auto_assign_mcp_server_to_team and
  _remove_mcp_server_from_team helpers for reuse and testability
- Update tests to mock at the correct boundaries

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../mcp_management_endpoints.py               | 119 ++++++++++--------
 .../test_mcp_manager_role.py                  |  89 ++++---------
 2 files changed, 96 insertions(+), 112 deletions(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index f137902f1b..e137ad15f1 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -203,6 +203,69 @@ if MCP_AVAILABLE:
 
         return resolved_team_id, team_obj
 
+    async def _auto_assign_mcp_server_to_team(
+        server_id: str,
+        team_id: str,
+        team_obj: "LiteLLM_TeamTableCachedObj",
+        prisma_client: Any,
+    ) -> None:
+        """
+        Add an MCP server to a team's ObjectPermissionTable and link the
+        permission back to the team if it didn't have one yet.
+
+        Uses handle_update_object_permission (the team-endpoint helper) so
+        the object_permission_id linkage follows the same pattern as
+        team_endpoints.update_team.
+        """
+        from litellm.proxy.management_endpoints.team_endpoints import (
+            handle_update_object_permission,
+        )
+
+        existing_mcp_servers: list = []
+        if team_obj.object_permission is not None:
+            existing_mcp_servers = team_obj.object_permission.mcp_servers or []
+        updated_mcp_servers = list(set(existing_mcp_servers + [server_id]))
+
+        # Build the data dict in the same shape team_endpoints uses
+        data_json: Dict[str, Any] = {
+            "object_permission": {"mcp_servers": updated_mcp_servers},
+        }
+        data_json = await handle_update_object_permission(
+            data_json=data_json,
+            existing_team_row=team_obj,
+        )
+
+        # If handle_update_object_permission produced an object_permission_id,
+        # persist it on the team row (it sets data_json["object_permission_id"]).
+        if "object_permission_id" in data_json:
+            await prisma_client.db.litellm_teamtable.update(
+                where={"team_id": team_id},
+                data={"object_permission_id": data_json["object_permission_id"]},
+            )
+
+    async def _remove_mcp_server_from_team(
+        server_id: str,
+        team_obj: "LiteLLM_TeamTableCachedObj",
+    ) -> None:
+        """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list."""
+        from litellm.proxy.proxy_server import prisma_client
+
+        existing_permission_id = getattr(team_obj, "object_permission_id", None)
+        if (
+            existing_permission_id is None
+            or team_obj.object_permission is None
+        ):
+            return
+
+        existing_mcp_servers = team_obj.object_permission.mcp_servers or []
+        updated_mcp_servers = [s for s in existing_mcp_servers if s != server_id]
+
+        await handle_update_object_permission_common(
+            data_json={"object_permission": {"mcp_servers": updated_mcp_servers}},
+            existing_object_permission_id=existing_permission_id,
+            prisma_client=prisma_client,
+        )
+
     @dataclass
     class _TemporaryMCPServerEntry:
         server: MCPServer
@@ -1347,36 +1410,12 @@ if MCP_AVAILABLE:
         # so a failure here doesn't mask the successfully created server).
         if manager_team_id is not None and manager_team_obj is not None:
             try:
-                existing_permission_id = getattr(
-                    manager_team_obj, "object_permission_id", None
-                )
-
-                # Read existing mcp_servers list and append the new server
-                existing_mcp_servers: list = []
-                if manager_team_obj.object_permission is not None:
-                    existing_mcp_servers = (
-                        manager_team_obj.object_permission.mcp_servers or []
-                    )
-                updated_mcp_servers = list(
-                    set(existing_mcp_servers + [new_mcp_server.server_id])
-                )
-
-                new_permission_id = await handle_update_object_permission_common(
-                    data_json={
-                        "object_permission": {
-                            "mcp_servers": updated_mcp_servers,
-                        }
-                    },
-                    existing_object_permission_id=existing_permission_id,
+                await _auto_assign_mcp_server_to_team(
+                    server_id=new_mcp_server.server_id,
+                    team_id=manager_team_id,
+                    team_obj=manager_team_obj,
                     prisma_client=prisma_client,
                 )
-
-                # If the team had no object_permission_id, link the new one
-                if existing_permission_id is None and new_permission_id is not None:
-                    await prisma_client.db.litellm_teamtable.update(
-                        where={"team_id": manager_team_id},
-                        data={"object_permission_id": new_permission_id},
-                    )
             except Exception as e:
                 verbose_proxy_logger.exception(
                     f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}"
@@ -1617,28 +1656,10 @@ if MCP_AVAILABLE:
         # Remove server from the manager's team permission list
         if manager_team_obj is not None:
             try:
-                existing_permission_id = getattr(
-                    manager_team_obj, "object_permission_id", None
+                await _remove_mcp_server_from_team(
+                    server_id=server_id,
+                    team_obj=manager_team_obj,
                 )
-                if (
-                    existing_permission_id is not None
-                    and manager_team_obj.object_permission is not None
-                ):
-                    existing_mcp_servers = (
-                        manager_team_obj.object_permission.mcp_servers or []
-                    )
-                    updated_mcp_servers = [
-                        s for s in existing_mcp_servers if s != server_id
-                    ]
-                    await handle_update_object_permission_common(
-                        data_json={
-                            "object_permission": {
-                                "mcp_servers": updated_mcp_servers,
-                            }
-                        },
-                        existing_object_permission_id=existing_permission_id,
-                        prisma_client=prisma_client,
-                    )
             except Exception as e:
                 verbose_proxy_logger.exception(
                     f"MCP server deleted but failed to remove from team permissions: {str(e)}"
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
index 67e9dda4ce..ed601dfffa 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -233,7 +233,7 @@ class TestCreateMcpServerAsManager:
         created_server.server_id = "new_server_id"
         created_server.credentials = None
 
-        mock_handle_update = AsyncMock()
+        mock_auto_assign = AsyncMock()
 
         with (
             patch(
@@ -263,38 +263,23 @@ class TestCreateMcpServerAsManager:
                 ),
             ),
             patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
-                mock_handle_update,
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._auto_assign_mcp_server_to_team",
+                mock_auto_assign,
             ),
         ):
             await add_mcp_server(payload=payload, user_api_key_dict=user_auth)
 
-            # Verify handle_update_object_permission_common was called with merged server list
-            mock_handle_update.assert_called_once()
-            call_kwargs = mock_handle_update.call_args.kwargs
-            mcp_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"]
-            assert "existing_server" in mcp_servers
-            assert "new_server_id" in mcp_servers
-            assert call_kwargs["existing_object_permission_id"] == "perm1"
+            # Verify _auto_assign_mcp_server_to_team was called with the right args
+            mock_auto_assign.assert_called_once()
+            call_kwargs = mock_auto_assign.call_args.kwargs
+            assert call_kwargs["server_id"] == "new_server_id"
+            assert call_kwargs["team_id"] == "team1"
+            assert call_kwargs["team_obj"] == mock_team
 
-    async def test_create_links_new_permission_to_team_when_none_exists(self):
-        """When team has no object_permission_id, create should link the new one."""
-        from litellm.proxy._types import NewMCPServerRequest
+    async def test_auto_assign_links_new_permission_to_team(self):
+        """_auto_assign_mcp_server_to_team should create permission and link to team."""
         from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            add_mcp_server,
-        )
-
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-            team_id="team1",
-        )
-
-        payload = NewMCPServerRequest(
-            server_name="test_server",
-            url="https://example.com/mcp",
-            team_id="team1",
+            _auto_assign_mcp_server_to_team,
         )
 
         # Team with NO object_permission_id
@@ -307,46 +292,24 @@ class TestCreateMcpServerAsManager:
         )
         mock_team.object_permission = None
 
-        created_server = MagicMock()
-        created_server.server_id = "new_server_id"
-        created_server.credentials = None
-
         mock_prisma = MagicMock()
         mock_prisma.db.litellm_teamtable.update = AsyncMock()
 
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
-                return_value=mock_prisma,
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(return_value=("team1", mock_team)),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
-                AsyncMock(return_value=None),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server",
-                AsyncMock(return_value=created_server),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
-                MagicMock(
-                    add_server=AsyncMock(),
-                    reload_servers_from_database=AsyncMock(),
-                ),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
-                AsyncMock(return_value="new_perm_id"),
-            ),
+        # handle_update_object_permission sets object_permission_id in data_json
+        async def fake_handle(data_json, existing_team_row):
+            data_json["object_permission_id"] = "new_perm_id"
+            return data_json
+
+        with patch(
+            "litellm.proxy.management_endpoints.team_endpoints.handle_update_object_permission",
+            side_effect=fake_handle,
         ):
-            await add_mcp_server(payload=payload, user_api_key_dict=user_auth)
+            await _auto_assign_mcp_server_to_team(
+                server_id="new_server_id",
+                team_id="team1",
+                team_obj=mock_team,
+                prisma_client=mock_prisma,
+            )
 
             # Verify the team was updated with the new object_permission_id
             mock_prisma.db.litellm_teamtable.update.assert_called_once_with(

From 1cd7a48c33cd72239913ccf73ecd87c3be4d5a30 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 09:19:53 -0700
Subject: [PATCH 23/31] Add tests for edit and delete MCP server manager paths

Addresses Greptile feedback about missing integration tests for PUT/DELETE
when invoked by mcp_server_manager role. Adds tests for edit success/403,
delete success with team cleanup/403, and the _remove_mcp_server_from_team
helper directly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../test_mcp_manager_role.py                  | 239 ++++++++++++++++++
 1 file changed, 239 insertions(+)

diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
index ed601dfffa..36cecde65b 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -69,6 +69,7 @@ class TestIsUserTeamMcpManager:
 
 
 from unittest.mock import AsyncMock, MagicMock, patch
+from fastapi import HTTPException
 from litellm.proxy._types import LiteLLM_TeamTableCachedObj
 
 
@@ -316,3 +317,241 @@ class TestCreateMcpServerAsManager:
                 where={"team_id": "team1"},
                 data={"object_permission_id": "new_perm_id"},
             )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
+)
+class TestEditMcpServerAsManager:
+    async def test_edit_succeeds_for_mcp_manager(self):
+        """MCP manager should be able to edit a server assigned to their team."""
+        from litellm.proxy._types import UpdateMCPServerRequest
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            edit_mcp_server,
+        )
+
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+            team_id="team1",
+        )
+
+        payload = UpdateMCPServerRequest(
+            server_id="server1",
+            description="Updated description",
+        )
+
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager"),
+            ],
+        )
+
+        updated_server = MagicMock()
+        updated_server.server_id = "server1"
+        updated_server.credentials = None
+
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+                AsyncMock(return_value=("team1", mock_team)),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server",
+                AsyncMock(return_value=updated_server),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+                MagicMock(
+                    update_server=AsyncMock(),
+                    reload_servers_from_database=AsyncMock(),
+                ),
+            ),
+        ):
+            result = await edit_mcp_server(payload=payload, user_api_key_dict=user_auth)
+            # Should not raise — edit succeeded
+            assert result is not None
+
+    async def test_edit_fails_for_server_not_in_team(self):
+        """MCP manager should get 403 when editing a server not in their team."""
+        from litellm.proxy._types import UpdateMCPServerRequest
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            edit_mcp_server,
+        )
+
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+            team_id="team1",
+        )
+
+        payload = UpdateMCPServerRequest(
+            server_id="server_not_in_team",
+            description="Updated description",
+        )
+
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+                AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")),
+            ),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await edit_mcp_server(payload=payload, user_api_key_dict=user_auth)
+            assert exc_info.value.status_code == 403
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
+)
+class TestDeleteMcpServerAsManager:
+    async def test_delete_succeeds_and_cleans_up_team(self):
+        """MCP manager deleting a server should also remove it from team permissions."""
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            remove_mcp_server,
+        )
+
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+            team_id="team1",
+        )
+
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager"),
+            ],
+        )
+
+        deleted_server = MagicMock()
+        deleted_server.server_id = "server1"
+
+        mock_remove_from_team = AsyncMock()
+
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+                AsyncMock(return_value=("team1", mock_team)),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_mcp_server",
+                AsyncMock(return_value=deleted_server),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
+                MagicMock(
+                    remove_server=MagicMock(),
+                    reload_servers_from_database=AsyncMock(),
+                ),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._remove_mcp_server_from_team",
+                mock_remove_from_team,
+            ),
+        ):
+            response = await remove_mcp_server(
+                server_id="server1", user_api_key_dict=user_auth
+            )
+            assert response.status_code == 202
+
+            # Verify team cleanup was called
+            mock_remove_from_team.assert_called_once()
+            call_kwargs = mock_remove_from_team.call_args.kwargs
+            assert call_kwargs["server_id"] == "server1"
+            assert call_kwargs["team_obj"] == mock_team
+
+    async def test_delete_fails_for_server_not_in_team(self):
+        """MCP manager should get 403 when deleting a server not in their team."""
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            remove_mcp_server,
+        )
+
+        user_auth = UserAPIKeyAuth(
+            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_id="user1",
+            api_key="sk-test",
+            team_id="team1",
+        )
+
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
+                return_value=MagicMock(),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+                AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")),
+            ),
+        ):
+            with pytest.raises(HTTPException) as exc_info:
+                await remove_mcp_server(
+                    server_id="server_not_in_team", user_api_key_dict=user_auth
+                )
+            assert exc_info.value.status_code == 403
+
+    async def test_remove_mcp_server_from_team_helper(self):
+        """_remove_mcp_server_from_team should update the permission list without the deleted server."""
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            _remove_mcp_server_from_team,
+        )
+
+        mock_team = LiteLLM_TeamTableCachedObj(
+            team_id="team1",
+            members_with_roles=[
+                Member(user_id="user1", role="mcp_server_manager"),
+            ],
+            object_permission_id="perm1",
+        )
+        mock_team.object_permission = MagicMock(
+            mcp_servers=["server1", "server2", "server3"]
+        )
+
+        mock_handle_common = AsyncMock()
+
+        with (
+            patch(
+                "litellm.proxy.proxy_server.prisma_client",
+                MagicMock(),
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
+                mock_handle_common,
+            ),
+        ):
+            await _remove_mcp_server_from_team(
+                server_id="server2",
+                team_obj=mock_team,
+            )
+
+            mock_handle_common.assert_called_once()
+            call_kwargs = mock_handle_common.call_args.kwargs
+            updated_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"]
+            assert "server2" not in updated_servers
+            assert "server1" in updated_servers
+            assert "server3" in updated_servers
+            assert call_kwargs["existing_object_permission_id"] == "perm1"

From 700fd86de971f0ca50bafafed129ab17af661f3a Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 12:13:09 -0700
Subject: [PATCH 24/31] Fix importorskip guard and add
 LiteLLM_TeamTableCachedObj import

- Add pytest.importorskip("mcp") at module level so tests skip cleanly
  in CI environments without the mcp package (instead of ImportError)
- Import LiteLLM_TeamTableCachedObj into MCP_AVAILABLE block so type
  annotations resolve for static analysis and get_type_hints()
- Remove string quotes from type annotations now that the import exists

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../proxy/management_endpoints/mcp_management_endpoints.py | 7 ++++---
 .../proxy/management_endpoints/test_mcp_manager_role.py    | 2 ++
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index e137ad15f1..8c8bd62988 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -112,6 +112,7 @@ if MCP_AVAILABLE:
     )
     from litellm.proxy._types import (
         LiteLLM_MCPServerTable,
+        LiteLLM_TeamTableCachedObj,
         LitellmUserRoles,
         MakeMCPServersPublicRequest,
         MCPApprovalStatus,
@@ -147,7 +148,7 @@ if MCP_AVAILABLE:
         user_api_key_dict: UserAPIKeyAuth,
         team_id: Optional[str] = None,
         server_id: Optional[str] = None,
-    ) -> Tuple[str, "LiteLLM_TeamTableCachedObj"]:
+    ) -> Tuple[str, LiteLLM_TeamTableCachedObj]:
         """
         Verify that the caller is an MCP server manager for a team and (for edit/delete)
         that the target server belongs to that team.
@@ -206,7 +207,7 @@ if MCP_AVAILABLE:
     async def _auto_assign_mcp_server_to_team(
         server_id: str,
         team_id: str,
-        team_obj: "LiteLLM_TeamTableCachedObj",
+        team_obj: LiteLLM_TeamTableCachedObj,
         prisma_client: Any,
     ) -> None:
         """
@@ -245,7 +246,7 @@ if MCP_AVAILABLE:
 
     async def _remove_mcp_server_from_team(
         server_id: str,
-        team_obj: "LiteLLM_TeamTableCachedObj",
+        team_obj: LiteLLM_TeamTableCachedObj,
     ) -> None:
         """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list."""
         from litellm.proxy.proxy_server import prisma_client
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
index 36cecde65b..86b0419fb2 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
@@ -1,5 +1,7 @@
 import pytest
 
+pytest.importorskip("mcp", reason="mcp package not installed")
+
 from litellm.proxy._types import (
     LiteLLM_TeamTable,
     LitellmUserRoles,

From 404c68c74b62a220fc935f2510fa649860f58aa4 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 12:42:35 -0700
Subject: [PATCH 25/31] Add missing team_member_budget_duration param to
 new_team docstring

Fixes CI failure in test_api_docs.py which validates that all Pydantic
model fields are documented in endpoint docstrings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 litellm/proxy/management_endpoints/team_endpoints.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 633de86aa6..23d4884592 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -701,6 +701,7 @@ async def new_team(  # noqa: PLR0915
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
     - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
     - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
+    - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
     - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
     - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
     - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"

From ba4aae02c75a4087e8041e6781c0501904ea6f1d Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 13:10:10 -0700
Subject: [PATCH 26/31] Fix outdated MCP server auth test for team MCP manager
 flow

The test_create_mcp_server_auth_failure test expected a 403 for non-admin
users, but the team MCP manager feature changed the auth flow to first
check for team_id (400) before checking permissions. Split into two tests:
one for missing team_id (400) and one for non-manager rejection (403).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../test_mcp_servers.py                       | 60 +++++++++++++++----
 1 file changed, 49 insertions(+), 11 deletions(-)

diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py
index a369cce83c..68ef977615 100644
--- a/tests/store_model_in_db_tests/test_mcp_servers.py
+++ b/tests/store_model_in_db_tests/test_mcp_servers.py
@@ -6,7 +6,7 @@ import os
 import asyncio
 from unittest import mock
 from fastapi.testclient import TestClient
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 
 from starlette import status
 
@@ -259,45 +259,83 @@ async def test_create_duplicate_mcp_server():
 
 
 @pytest.mark.asyncio
-async def test_create_mcp_server_auth_failure():
+async def test_create_mcp_server_auth_failure_no_team_id():
     """
-    Test that non-admin users cannot create MCP servers.
+    Test that non-admin users without a team_id get a 400 error
+    requiring team_id for team MCP manager flow.
     """
-    # Mock the database functions directly
     with mock.patch(
         "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE",
         True,
     ), mock.patch(
         "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw"
     ) as mock_get_prisma:
-        # Import after mocking
         from litellm.proxy.management_endpoints.mcp_management_endpoints import (
             add_mcp_server,
         )
         from fastapi import HTTPException
 
-        # Mock database client
         mock_prisma = mock.Mock()
         mock_get_prisma.return_value = mock_prisma
 
-        # Set up test data
         server_id = str(uuid.uuid4())
         mcp_server_request = generate_mcpserver_create_request(server_id=server_id)
 
-        # Create mock user auth without admin role
         user_auth = UserAPIKeyAuth(
             api_key=TEST_MASTER_KEY,
             user_id="test-user",
-            user_role=LitellmUserRoles.INTERNAL_USER,  # Not an admin
+            user_role=LitellmUserRoles.INTERNAL_USER,
+        )
+
+        with pytest.raises(HTTPException) as exc_info:
+            await add_mcp_server(
+                payload=mcp_server_request, user_api_key_dict=user_auth
+            )
+
+        assert exc_info.value.status_code == 400
+        assert "team_id is required" in str(exc_info.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_create_mcp_server_auth_failure_not_manager():
+    """
+    Test that non-admin users with a team_id but without MCP manager
+    permissions get a 403 error.
+    """
+    with mock.patch(
+        "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE",
+        True,
+    ), mock.patch(
+        "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw"
+    ) as mock_get_prisma, mock.patch(
+        "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
+        side_effect=HTTPException(
+            status_code=403,
+            detail={"error": "You do not have permission to manage MCP servers for this team."},
+        ),
+    ):
+        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
+            add_mcp_server,
+        )
+
+        mock_prisma = mock.Mock()
+        mock_get_prisma.return_value = mock_prisma
+
+        server_id = str(uuid.uuid4())
+        mcp_server_request = generate_mcpserver_create_request(server_id=server_id)
+        mcp_server_request.team_id = "some-team-id"
+
+        user_auth = UserAPIKeyAuth(
+            api_key=TEST_MASTER_KEY,
+            user_id="test-user",
+            user_role=LitellmUserRoles.INTERNAL_USER,
         )
 
-        # Expect HTTPException to be raised
         with pytest.raises(HTTPException) as exc_info:
             await add_mcp_server(
                 payload=mcp_server_request, user_api_key_dict=user_auth
             )
 
-        # Verify the exception details
         assert exc_info.value.status_code == 403
         assert "permission" in str(exc_info.value.detail)
 

From 6f81eb01fe1b1efad07b801d21973a30c8aa68b0 Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <rcrabbe@berkeley.edu>
Date: Fri, 20 Mar 2026 13:57:26 -0700
Subject: [PATCH 27/31] fix: apply Black formatting to ui_sso.py

---
 litellm/proxy/management_endpoints/ui_sso.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
index 8634a8bb33..d06ce56f81 100644
--- a/litellm/proxy/management_endpoints/ui_sso.py
+++ b/litellm/proxy/management_endpoints/ui_sso.py
@@ -2606,9 +2606,7 @@ class SSOAuthenticationHandler:
 
             separator = "&" if "?" in return_to else "?"
             redirect_url = (
-                return_to
-                + separator
-                + urlencode({"login": "success", "code": code})
+                return_to + separator + urlencode({"login": "success", "code": code})
             )
             verbose_proxy_logger.info(
                 "Cross-origin SSO: redirecting to control plane with login code"

From c9683c6f9751b6877b938a5b2f29937394c8179d Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 15:41:57 -0700
Subject: [PATCH 28/31] Revert "[Feature] Team MCP Server Manager Role"

---
 litellm/proxy/_experimental/mcp_server/db.py  |   4 +-
 litellm/proxy/_types.py                       |   7 +-
 .../management_endpoints/common_utils.py      |  12 -
 .../mcp_management_endpoints.py               | 221 +------
 .../test_mcp_manager_role.py                  | 559 ------------------
 5 files changed, 31 insertions(+), 772 deletions(-)
 delete mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py

diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py
index 0bc102889b..fbef33c32e 100644
--- a/litellm/proxy/_experimental/mcp_server/db.py
+++ b/litellm/proxy/_experimental/mcp_server/db.py
@@ -40,8 +40,8 @@ def _prepare_mcp_server_data(
     """
     from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
-    # Convert model to dict, excluding fields not in the DB schema
-    data_dict = data.model_dump(exclude_none=True, exclude={"team_id"})
+    # Convert model to dict
+    data_dict = data.model_dump(exclude_none=True)
     # Ensure alias is always present in the dict (even if None)
     if "alias" not in data_dict:
         data_dict["alias"] = getattr(data, "alias", None)
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 826e5c9a7b..91a953c217 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1144,10 +1144,6 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
         None,
         description="Server-managed: set by the endpoint; caller values are overridden.",
     )
-    team_id: Optional[str] = Field(
-        None,
-        description="Team ID to assign the MCP server to. Required for team MCP managers.",
-    )
 
     @model_validator(mode="before")
     @classmethod
@@ -1612,9 +1608,8 @@ class Member(MemberBase):
     role: Literal[
         "admin",
         "user",
-        "mcp_server_manager",
     ] = Field(
-        description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member, 'mcp_server_manager' can manage MCP servers for the team"
+        description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member"
     )
 
 
diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py
index 092b22cee6..efc42d3355 100644
--- a/litellm/proxy/management_endpoints/common_utils.py
+++ b/litellm/proxy/management_endpoints/common_utils.py
@@ -41,18 +41,6 @@ def _is_user_team_admin(
     return False
 
 
-def _is_user_team_mcp_manager(
-    user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
-) -> bool:
-    for member in team_obj.members_with_roles:
-        if (
-            member.user_id is not None and member.user_id == user_api_key_dict.user_id
-        ) and member.role == "mcp_server_manager":
-            return True
-
-    return False
-
-
 async def _is_user_org_admin_for_team(
     user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
 ) -> bool:
diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index 883e1a94aa..f29a721ede 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -21,7 +21,7 @@ import json
 import os
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
+from typing import Any, Dict, Iterable, List, Literal, Optional
 
 from fastapi import (
     APIRouter,
@@ -112,7 +112,6 @@ if MCP_AVAILABLE:
     )
     from litellm.proxy._types import (
         LiteLLM_MCPServerTable,
-        LiteLLM_TeamTableCachedObj,
         LitellmUserRoles,
         MakeMCPServersPublicRequest,
         MCPApprovalStatus,
@@ -131,142 +130,11 @@ if MCP_AVAILABLE:
     )
     from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
     from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
-    from litellm.proxy.auth.auth_checks import get_team_object
-    from litellm.proxy.management_endpoints.common_utils import (
-        _is_user_team_mcp_manager,
-        _user_has_admin_view,
-    )
-    from litellm.proxy.management_helpers.object_permission_utils import (
-        _get_team_allowed_mcp_servers,
-        handle_update_object_permission_common,
-    )
+    from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
     from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
     from litellm.types.mcp import MCPCredentials
     from litellm.types.mcp_server.mcp_server_manager import MCPServer
 
-    async def _assert_can_manage_team_mcp_server(
-        user_api_key_dict: UserAPIKeyAuth,
-        team_id: Optional[str] = None,
-        server_id: Optional[str] = None,
-    ) -> Tuple[str, LiteLLM_TeamTableCachedObj]:
-        """
-        Verify that the caller is an MCP server manager for a team and (for edit/delete)
-        that the target server belongs to that team.
-
-        Returns a tuple of (team_id, team_obj) for downstream use.
-        Raises HTTPException(400) if no team_id can be determined.
-        Raises HTTPException(403) if the caller is not an MCP manager or server not in team.
-        """
-        from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
-
-        resolved_team_id = team_id or user_api_key_dict.team_id
-        if not resolved_team_id:
-            raise HTTPException(
-                status_code=400,
-                detail={"error": "team_id is required for MCP server manager operations."},
-            )
-
-        # When the API key is team-scoped, ensure the request team_id matches
-        if (
-            team_id
-            and user_api_key_dict.team_id
-            and team_id != user_api_key_dict.team_id
-        ):
-            raise HTTPException(
-                status_code=403,
-                detail={"error": "team_id does not match the API key's team."},
-            )
-
-        team_obj = await get_team_object(
-            team_id=resolved_team_id,
-            prisma_client=prisma_client,
-            user_api_key_cache=user_api_key_cache,
-            check_db_only=True,  # bypass cache to get fresh object_permission
-        )
-
-        if not _is_user_team_mcp_manager(user_api_key_dict, team_obj):
-            raise HTTPException(
-                status_code=403,
-                detail={
-                    "error": f"User does not have mcp_server_manager role in team {resolved_team_id}."
-                },
-            )
-
-        if server_id is not None:
-            team_server_ids = await _get_team_allowed_mcp_servers(team_obj)
-            if server_id not in team_server_ids:
-                raise HTTPException(
-                    status_code=403,
-                    detail={
-                        "error": f"MCP server {server_id} is not assigned to team {resolved_team_id}."
-                    },
-                )
-
-        return resolved_team_id, team_obj
-
-    async def _auto_assign_mcp_server_to_team(
-        server_id: str,
-        team_id: str,
-        team_obj: LiteLLM_TeamTableCachedObj,
-        prisma_client: Any,
-    ) -> None:
-        """
-        Add an MCP server to a team's ObjectPermissionTable and link the
-        permission back to the team if it didn't have one yet.
-
-        Uses handle_update_object_permission (the team-endpoint helper) so
-        the object_permission_id linkage follows the same pattern as
-        team_endpoints.update_team.
-        """
-        from litellm.proxy.management_endpoints.team_endpoints import (
-            handle_update_object_permission,
-        )
-
-        existing_mcp_servers: list = []
-        if team_obj.object_permission is not None:
-            existing_mcp_servers = team_obj.object_permission.mcp_servers or []
-        updated_mcp_servers = list(set(existing_mcp_servers + [server_id]))
-
-        # Build the data dict in the same shape team_endpoints uses
-        data_json: Dict[str, Any] = {
-            "object_permission": {"mcp_servers": updated_mcp_servers},
-        }
-        data_json = await handle_update_object_permission(
-            data_json=data_json,
-            existing_team_row=team_obj,
-        )
-
-        # If handle_update_object_permission produced an object_permission_id,
-        # persist it on the team row (it sets data_json["object_permission_id"]).
-        if "object_permission_id" in data_json:
-            await prisma_client.db.litellm_teamtable.update(
-                where={"team_id": team_id},
-                data={"object_permission_id": data_json["object_permission_id"]},
-            )
-
-    async def _remove_mcp_server_from_team(
-        server_id: str,
-        team_obj: LiteLLM_TeamTableCachedObj,
-    ) -> None:
-        """Remove a server ID from a team's ObjectPermissionTable.mcp_servers list."""
-        from litellm.proxy.proxy_server import prisma_client
-
-        existing_permission_id = getattr(team_obj, "object_permission_id", None)
-        if (
-            existing_permission_id is None
-            or team_obj.object_permission is None
-        ):
-            return
-
-        existing_mcp_servers = team_obj.object_permission.mcp_servers or []
-        updated_mcp_servers = [s for s in existing_mcp_servers if s != server_id]
-
-        await handle_update_object_permission_common(
-            data_json={"object_permission": {"mcp_servers": updated_mcp_servers}},
-            existing_object_permission_id=existing_permission_id,
-            prisma_client=prisma_client,
-        )
-
     @dataclass
     class _TemporaryMCPServerEntry:
         server: MCPServer
@@ -1338,27 +1206,16 @@ if MCP_AVAILABLE:
         # Validate and normalize payload fields
         validate_and_normalize_mcp_server_payload(payload)
 
-        # AuthZ - proxy admins or team MCP managers can create MCP servers
-        is_proxy_admin = LitellmUserRoles.PROXY_ADMIN == user_api_key_dict.user_role
-        manager_team_id: Optional[str] = None
-
-        manager_team_obj = None
-        if not is_proxy_admin:
-            # Check if the user is an MCP manager for the specified team
-            if payload.team_id is None:
-                raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail={
-                        "error": "team_id is required when creating MCP servers as a team MCP manager."
-                    },
-                )
-            manager_team_id, manager_team_obj = await _assert_can_manage_team_mcp_server(
-                user_api_key_dict=user_api_key_dict,
-                team_id=payload.team_id,
+        # AuthZ - restrict only proxy admins to create mcp servers
+        if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail={
+                    "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN."
+                },
             )
-
-        # Fail if the MCP server with this id already exists
-        if payload.server_id is not None:
+        elif payload.server_id is not None:
+            # fail if the mcp server with id already exists
             mcp_server = await get_mcp_server(prisma_client, payload.server_id)
             if mcp_server is not None:
                 raise HTTPException(
@@ -1367,8 +1224,7 @@ if MCP_AVAILABLE:
                         "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
                     },
                 )
-
-        if (
+        elif (
             SpecialMCPServerName.all_team_servers == payload.server_id
             or SpecialMCPServerName.all_proxy_servers == payload.server_id
         ):
@@ -1399,28 +1255,12 @@ if MCP_AVAILABLE:
 
             # Ensure registry is up to date by reloading from database
             await global_mcp_server_manager.reload_servers_from_database()
-
         except Exception as e:
             verbose_proxy_logger.exception(f"Error creating mcp server: {str(e)}")
             raise HTTPException(
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 detail={"error": f"Error creating mcp server: {str(e)}"},
             )
-
-        # Auto-assign the server to the manager's team (separate from create
-        # so a failure here doesn't mask the successfully created server).
-        if manager_team_id is not None and manager_team_obj is not None:
-            try:
-                await _auto_assign_mcp_server_to_team(
-                    server_id=new_mcp_server.server_id,
-                    team_id=manager_team_id,
-                    team_obj=manager_team_obj,
-                    prisma_client=prisma_client,
-                )
-            except Exception as e:
-                verbose_proxy_logger.exception(
-                    f"MCP server created but failed to auto-assign to team {manager_team_id}: {str(e)}"
-                )
         return _redact_mcp_credentials(new_mcp_server)
 
     @router.post(
@@ -1637,12 +1477,15 @@ if MCP_AVAILABLE:
             "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
         )
 
-        # Authz - proxy admins or team MCP managers can delete MCP servers
-        manager_team_obj = None
+        # Authz - restrict only admins to delete mcp servers
         if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
-            _, manager_team_obj = await _assert_can_manage_team_mcp_server(
-                user_api_key_dict=user_api_key_dict,
-                server_id=server_id,
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail={
+                    "error": "Call not allowed to delete MCP server. User is not a proxy admin. route={}".format(
+                        "DELETE /v1/mcp/server"
+                    )
+                },
             )
 
         # try to delete the mcp server
@@ -1658,18 +1501,6 @@ if MCP_AVAILABLE:
         # Ensure registry is up to date by reloading from database
         await global_mcp_server_manager.reload_servers_from_database()
 
-        # Remove server from the manager's team permission list
-        if manager_team_obj is not None:
-            try:
-                await _remove_mcp_server_from_team(
-                    server_id=server_id,
-                    team_obj=manager_team_obj,
-                )
-            except Exception as e:
-                verbose_proxy_logger.exception(
-                    f"MCP server deleted but failed to remove from team permissions: {str(e)}"
-                )
-
         # TODO: Enterprise: Finish audit log trail
         if litellm.store_audit_logs:
             pass
@@ -1971,11 +1802,15 @@ if MCP_AVAILABLE:
         # Validate and normalize payload fields
         validate_and_normalize_mcp_server_payload(payload)
 
-        # Authz - proxy admins or team MCP managers can update MCP servers
+        # Authz - restrict only admins to delete mcp servers
         if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role:
-            await _assert_can_manage_team_mcp_server(
-                user_api_key_dict=user_api_key_dict,
-                server_id=payload.server_id,
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail={
+                    "error": "Call not allowed to update MCP server. User is not a proxy admin. route={}".format(
+                        "PUT /v1/mcp/server"
+                    )
+                },
             )
 
         # try to update the mcp server
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
deleted file mode 100644
index 86b0419fb2..0000000000
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py
+++ /dev/null
@@ -1,559 +0,0 @@
-import pytest
-
-pytest.importorskip("mcp", reason="mcp package not installed")
-
-from litellm.proxy._types import (
-    LiteLLM_TeamTable,
-    LitellmUserRoles,
-    Member,
-    UserAPIKeyAuth,
-)
-from litellm.proxy.management_endpoints.common_utils import (
-    _is_user_team_mcp_manager,
-)
-from litellm.proxy.management_endpoints import (
-    mcp_management_endpoints as mgmt_endpoints,
-)
-
-
-class TestIsUserTeamMcpManager:
-    def test_mcp_server_manager_role_returns_true(self):
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-        )
-        team = LiteLLM_TeamTable(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager")
-            ],
-        )
-        assert _is_user_team_mcp_manager(user_auth, team) is True
-
-    def test_regular_user_role_returns_false(self):
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-        )
-        team = LiteLLM_TeamTable(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="user")],
-        )
-        assert _is_user_team_mcp_manager(user_auth, team) is False
-
-    def test_admin_role_returns_false(self):
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-        )
-        team = LiteLLM_TeamTable(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="admin")],
-        )
-        assert _is_user_team_mcp_manager(user_auth, team) is False
-
-    def test_user_not_in_team_returns_false(self):
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user2",
-            api_key="sk-test",
-        )
-        team = LiteLLM_TeamTable(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager")
-            ],
-        )
-        assert _is_user_team_mcp_manager(user_auth, team) is False
-
-
-from unittest.mock import AsyncMock, MagicMock, patch
-from fastapi import HTTPException
-from litellm.proxy._types import LiteLLM_TeamTableCachedObj
-
-
-@pytest.mark.asyncio
-@pytest.mark.skipif(
-    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
-)
-class TestAssertCanManageTeamMcpServer:
-    async def test_mcp_manager_with_team_id_succeeds(self):
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _assert_can_manage_team_mcp_server,
-        )
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test",
-        )
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="mcp_server_manager")],
-        )
-        with patch(
-            "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
-            AsyncMock(return_value=mock_team),
-        ):
-            team_id, team_obj = await _assert_can_manage_team_mcp_server(
-                user_api_key_dict=user_auth, team_id="team1"
-            )
-            assert team_id == "team1"
-            assert team_obj == mock_team
-
-    async def test_regular_user_gets_403(self):
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _assert_can_manage_team_mcp_server,
-        )
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test",
-        )
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="user")],
-        )
-        with patch(
-            "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
-            AsyncMock(return_value=mock_team),
-        ):
-            with pytest.raises(Exception) as exc_info:
-                await _assert_can_manage_team_mcp_server(
-                    user_api_key_dict=user_auth, team_id="team1"
-                )
-            assert exc_info.value.status_code == 403
-
-    async def test_admin_gets_403(self):
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _assert_can_manage_team_mcp_server,
-        )
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test",
-        )
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="admin")],
-        )
-        with patch(
-            "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
-            AsyncMock(return_value=mock_team),
-        ):
-            with pytest.raises(Exception) as exc_info:
-                await _assert_can_manage_team_mcp_server(
-                    user_api_key_dict=user_auth, team_id="team1"
-                )
-            assert exc_info.value.status_code == 403
-
-    async def test_mcp_manager_server_in_team_succeeds(self):
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _assert_can_manage_team_mcp_server,
-        )
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1",
-        )
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="mcp_server_manager")],
-        )
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
-                AsyncMock(return_value=mock_team),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers",
-                AsyncMock(return_value={"server1", "server2"}),
-            ),
-        ):
-            team_id, team_obj = await _assert_can_manage_team_mcp_server(
-                user_api_key_dict=user_auth, server_id="server1"
-            )
-            assert team_id == "team1"
-            assert team_obj == mock_team
-
-    async def test_mcp_manager_server_not_in_team_gets_403(self):
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _assert_can_manage_team_mcp_server,
-        )
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER, user_id="user1", api_key="sk-test", team_id="team1",
-        )
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[Member(user_id="user1", role="mcp_server_manager")],
-        )
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object",
-                AsyncMock(return_value=mock_team),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_allowed_mcp_servers",
-                AsyncMock(return_value={"server2", "server3"}),
-            ),
-        ):
-            with pytest.raises(Exception) as exc_info:
-                await _assert_can_manage_team_mcp_server(
-                    user_api_key_dict=user_auth, server_id="server1"
-                )
-            assert exc_info.value.status_code == 403
-
-
-@pytest.mark.asyncio
-@pytest.mark.skipif(
-    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
-)
-class TestCreateMcpServerAsManager:
-    async def test_create_auto_assigns_to_team(self):
-        """MCP manager creating a server should auto-assign it to their team's ObjectPermissionTable."""
-        from litellm.proxy._types import NewMCPServerRequest
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            add_mcp_server,
-        )
-
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-            team_id="team1",
-        )
-
-        payload = NewMCPServerRequest(
-            server_name="test-server",
-            url="https://example.com/mcp",
-            team_id="team1",
-        )
-
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager"),
-            ],
-            object_permission_id="perm1",
-        )
-        mock_team.object_permission = MagicMock(mcp_servers=["existing_server"])
-
-        created_server = MagicMock()
-        created_server.server_id = "new_server_id"
-        created_server.credentials = None
-
-        mock_auto_assign = AsyncMock()
-
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(return_value=("team1", mock_team)),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server",
-                AsyncMock(return_value=None),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server",
-                AsyncMock(return_value=created_server),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
-                MagicMock(
-                    add_server=AsyncMock(),
-                    reload_servers_from_database=AsyncMock(),
-                ),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._auto_assign_mcp_server_to_team",
-                mock_auto_assign,
-            ),
-        ):
-            await add_mcp_server(payload=payload, user_api_key_dict=user_auth)
-
-            # Verify _auto_assign_mcp_server_to_team was called with the right args
-            mock_auto_assign.assert_called_once()
-            call_kwargs = mock_auto_assign.call_args.kwargs
-            assert call_kwargs["server_id"] == "new_server_id"
-            assert call_kwargs["team_id"] == "team1"
-            assert call_kwargs["team_obj"] == mock_team
-
-    async def test_auto_assign_links_new_permission_to_team(self):
-        """_auto_assign_mcp_server_to_team should create permission and link to team."""
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _auto_assign_mcp_server_to_team,
-        )
-
-        # Team with NO object_permission_id
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager"),
-            ],
-            object_permission_id=None,
-        )
-        mock_team.object_permission = None
-
-        mock_prisma = MagicMock()
-        mock_prisma.db.litellm_teamtable.update = AsyncMock()
-
-        # handle_update_object_permission sets object_permission_id in data_json
-        async def fake_handle(data_json, existing_team_row):
-            data_json["object_permission_id"] = "new_perm_id"
-            return data_json
-
-        with patch(
-            "litellm.proxy.management_endpoints.team_endpoints.handle_update_object_permission",
-            side_effect=fake_handle,
-        ):
-            await _auto_assign_mcp_server_to_team(
-                server_id="new_server_id",
-                team_id="team1",
-                team_obj=mock_team,
-                prisma_client=mock_prisma,
-            )
-
-            # Verify the team was updated with the new object_permission_id
-            mock_prisma.db.litellm_teamtable.update.assert_called_once_with(
-                where={"team_id": "team1"},
-                data={"object_permission_id": "new_perm_id"},
-            )
-
-
-@pytest.mark.asyncio
-@pytest.mark.skipif(
-    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
-)
-class TestEditMcpServerAsManager:
-    async def test_edit_succeeds_for_mcp_manager(self):
-        """MCP manager should be able to edit a server assigned to their team."""
-        from litellm.proxy._types import UpdateMCPServerRequest
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            edit_mcp_server,
-        )
-
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-            team_id="team1",
-        )
-
-        payload = UpdateMCPServerRequest(
-            server_id="server1",
-            description="Updated description",
-        )
-
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager"),
-            ],
-        )
-
-        updated_server = MagicMock()
-        updated_server.server_id = "server1"
-        updated_server.credentials = None
-
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(return_value=("team1", mock_team)),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server",
-                AsyncMock(return_value=updated_server),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
-                MagicMock(
-                    update_server=AsyncMock(),
-                    reload_servers_from_database=AsyncMock(),
-                ),
-            ),
-        ):
-            result = await edit_mcp_server(payload=payload, user_api_key_dict=user_auth)
-            # Should not raise — edit succeeded
-            assert result is not None
-
-    async def test_edit_fails_for_server_not_in_team(self):
-        """MCP manager should get 403 when editing a server not in their team."""
-        from litellm.proxy._types import UpdateMCPServerRequest
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            edit_mcp_server,
-        )
-
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-            team_id="team1",
-        )
-
-        payload = UpdateMCPServerRequest(
-            server_id="server_not_in_team",
-            description="Updated description",
-        )
-
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload",
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")),
-            ),
-        ):
-            with pytest.raises(HTTPException) as exc_info:
-                await edit_mcp_server(payload=payload, user_api_key_dict=user_auth)
-            assert exc_info.value.status_code == 403
-
-
-@pytest.mark.asyncio
-@pytest.mark.skipif(
-    not mgmt_endpoints.MCP_AVAILABLE, reason="MCP module not installed"
-)
-class TestDeleteMcpServerAsManager:
-    async def test_delete_succeeds_and_cleans_up_team(self):
-        """MCP manager deleting a server should also remove it from team permissions."""
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            remove_mcp_server,
-        )
-
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-            team_id="team1",
-        )
-
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager"),
-            ],
-        )
-
-        deleted_server = MagicMock()
-        deleted_server.server_id = "server1"
-
-        mock_remove_from_team = AsyncMock()
-
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(return_value=("team1", mock_team)),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_mcp_server",
-                AsyncMock(return_value=deleted_server),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
-                MagicMock(
-                    remove_server=MagicMock(),
-                    reload_servers_from_database=AsyncMock(),
-                ),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._remove_mcp_server_from_team",
-                mock_remove_from_team,
-            ),
-        ):
-            response = await remove_mcp_server(
-                server_id="server1", user_api_key_dict=user_auth
-            )
-            assert response.status_code == 202
-
-            # Verify team cleanup was called
-            mock_remove_from_team.assert_called_once()
-            call_kwargs = mock_remove_from_team.call_args.kwargs
-            assert call_kwargs["server_id"] == "server1"
-            assert call_kwargs["team_obj"] == mock_team
-
-    async def test_delete_fails_for_server_not_in_team(self):
-        """MCP manager should get 403 when deleting a server not in their team."""
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            remove_mcp_server,
-        )
-
-        user_auth = UserAPIKeyAuth(
-            user_role=LitellmUserRoles.INTERNAL_USER,
-            user_id="user1",
-            api_key="sk-test",
-            team_id="team1",
-        )
-
-        with (
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
-                return_value=MagicMock(),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-                AsyncMock(side_effect=HTTPException(status_code=403, detail="Not in team")),
-            ),
-        ):
-            with pytest.raises(HTTPException) as exc_info:
-                await remove_mcp_server(
-                    server_id="server_not_in_team", user_api_key_dict=user_auth
-                )
-            assert exc_info.value.status_code == 403
-
-    async def test_remove_mcp_server_from_team_helper(self):
-        """_remove_mcp_server_from_team should update the permission list without the deleted server."""
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            _remove_mcp_server_from_team,
-        )
-
-        mock_team = LiteLLM_TeamTableCachedObj(
-            team_id="team1",
-            members_with_roles=[
-                Member(user_id="user1", role="mcp_server_manager"),
-            ],
-            object_permission_id="perm1",
-        )
-        mock_team.object_permission = MagicMock(
-            mcp_servers=["server1", "server2", "server3"]
-        )
-
-        mock_handle_common = AsyncMock()
-
-        with (
-            patch(
-                "litellm.proxy.proxy_server.prisma_client",
-                MagicMock(),
-            ),
-            patch(
-                "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common",
-                mock_handle_common,
-            ),
-        ):
-            await _remove_mcp_server_from_team(
-                server_id="server2",
-                team_obj=mock_team,
-            )
-
-            mock_handle_common.assert_called_once()
-            call_kwargs = mock_handle_common.call_args.kwargs
-            updated_servers = call_kwargs["data_json"]["object_permission"]["mcp_servers"]
-            assert "server2" not in updated_servers
-            assert "server1" in updated_servers
-            assert "server3" in updated_servers
-            assert call_kwargs["existing_object_permission_id"] == "perm1"

From 6862930538ed74ecd9411b24b11fd5928af3fa5a Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 15:44:00 -0700
Subject: [PATCH 29/31] Revert test to match reverted team MCP manager feature

The team MCP manager feature was reverted in PR #24255, so the test
needs to go back to the original single auth failure test that expects
a 403 for non-admin users.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../test_mcp_servers.py                       | 60 ++++---------------
 1 file changed, 11 insertions(+), 49 deletions(-)

diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py
index 68ef977615..a369cce83c 100644
--- a/tests/store_model_in_db_tests/test_mcp_servers.py
+++ b/tests/store_model_in_db_tests/test_mcp_servers.py
@@ -6,7 +6,7 @@ import os
 import asyncio
 from unittest import mock
 from fastapi.testclient import TestClient
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 
 from starlette import status
 
@@ -259,83 +259,45 @@ async def test_create_duplicate_mcp_server():
 
 
 @pytest.mark.asyncio
-async def test_create_mcp_server_auth_failure_no_team_id():
+async def test_create_mcp_server_auth_failure():
     """
-    Test that non-admin users without a team_id get a 400 error
-    requiring team_id for team MCP manager flow.
+    Test that non-admin users cannot create MCP servers.
     """
+    # Mock the database functions directly
     with mock.patch(
         "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE",
         True,
     ), mock.patch(
         "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw"
     ) as mock_get_prisma:
+        # Import after mocking
         from litellm.proxy.management_endpoints.mcp_management_endpoints import (
             add_mcp_server,
         )
         from fastapi import HTTPException
 
+        # Mock database client
         mock_prisma = mock.Mock()
         mock_get_prisma.return_value = mock_prisma
 
+        # Set up test data
         server_id = str(uuid.uuid4())
         mcp_server_request = generate_mcpserver_create_request(server_id=server_id)
 
+        # Create mock user auth without admin role
         user_auth = UserAPIKeyAuth(
             api_key=TEST_MASTER_KEY,
             user_id="test-user",
-            user_role=LitellmUserRoles.INTERNAL_USER,
-        )
-
-        with pytest.raises(HTTPException) as exc_info:
-            await add_mcp_server(
-                payload=mcp_server_request, user_api_key_dict=user_auth
-            )
-
-        assert exc_info.value.status_code == 400
-        assert "team_id is required" in str(exc_info.value.detail)
-
-
-@pytest.mark.asyncio
-async def test_create_mcp_server_auth_failure_not_manager():
-    """
-    Test that non-admin users with a team_id but without MCP manager
-    permissions get a 403 error.
-    """
-    with mock.patch(
-        "litellm.proxy.management_endpoints.mcp_management_endpoints.MCP_AVAILABLE",
-        True,
-    ), mock.patch(
-        "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw"
-    ) as mock_get_prisma, mock.patch(
-        "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server",
-        side_effect=HTTPException(
-            status_code=403,
-            detail={"error": "You do not have permission to manage MCP servers for this team."},
-        ),
-    ):
-        from litellm.proxy.management_endpoints.mcp_management_endpoints import (
-            add_mcp_server,
-        )
-
-        mock_prisma = mock.Mock()
-        mock_get_prisma.return_value = mock_prisma
-
-        server_id = str(uuid.uuid4())
-        mcp_server_request = generate_mcpserver_create_request(server_id=server_id)
-        mcp_server_request.team_id = "some-team-id"
-
-        user_auth = UserAPIKeyAuth(
-            api_key=TEST_MASTER_KEY,
-            user_id="test-user",
-            user_role=LitellmUserRoles.INTERNAL_USER,
+            user_role=LitellmUserRoles.INTERNAL_USER,  # Not an admin
         )
 
+        # Expect HTTPException to be raised
         with pytest.raises(HTTPException) as exc_info:
             await add_mcp_server(
                 payload=mcp_server_request, user_api_key_dict=user_auth
             )
 
+        # Verify the exception details
         assert exc_info.value.status_code == 403
         assert "permission" in str(exc_info.value.detail)
 

From f0e0d98f861bb5e42a0828618bdd75e2a7fc2514 Mon Sep 17 00:00:00 2001
From: joereyna <joseph.reyna@gmail.com>
Date: Fri, 20 Mar 2026 16:07:09 -0700
Subject: [PATCH 30/31] fix(test): mock get_auth_header instead of get_api_key
 in anthropic file content test

---
 .../llms/anthropic/test_anthropic_files_and_batches.py          | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py
index 4757122017..0755c189af 100644
--- a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py
+++ b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py
@@ -336,7 +336,7 @@ class TestAnthropicFilesHandler:
             "extra_body": None
         }
 
-        with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None):
+        with patch.object(handler.anthropic_model_info, "get_auth_header", return_value=None):
             with pytest.raises(ValueError, match="Missing Anthropic API Key"):
                 await handler.afile_content(
                     file_content_request=file_content_request,

From e678ddea4310754a001659e66c170439562cf2ae Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 20 Mar 2026 16:12:48 -0700
Subject: [PATCH 31/31] Fix unreachable special MCP server name guard in
 add_mcp_server

The special name check (all_team_servers, all_proxy_servers) was an elif
after the server_id-is-not-None check, making it unreachable since special
names are non-None strings. Split into separate if blocks so the special
name guard runs before the duplicate-ID check.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../mcp_management_endpoints.py               | 25 +++++++++++--------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
index f29a721ede..e4bb288cda 100644
--- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py
@@ -1214,17 +1214,9 @@ if MCP_AVAILABLE:
                     "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN."
                 },
             )
-        elif payload.server_id is not None:
-            # fail if the mcp server with id already exists
-            mcp_server = await get_mcp_server(prisma_client, payload.server_id)
-            if mcp_server is not None:
-                raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail={
-                        "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
-                    },
-                )
-        elif (
+
+        # Block reserved special server IDs
+        if (
             SpecialMCPServerName.all_team_servers == payload.server_id
             or SpecialMCPServerName.all_proxy_servers == payload.server_id
         ):
@@ -1235,6 +1227,17 @@ if MCP_AVAILABLE:
                 },
             )
 
+        if payload.server_id is not None:
+            # fail if the mcp server with id already exists
+            mcp_server = await get_mcp_server(prisma_client, payload.server_id)
+            if mcp_server is not None:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={
+                        "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
+                    },
+                )
+
         # TODO: audit log for create
 
         # Admin-created servers are always active — clear any submission lifecycle