Team Settings Model Select

2026-07-15 22:22:12 +00:00 · 2026-01-15 21:57:11 -08:00
parent 45d38bea5b
commit f81aefd48a
8 changed files with 505 additions and 204 deletions
@@ -18,6 +18,10 @@ vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({
  useOrganization: vi.fn(),
 }));

+vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({
+  useCurrentUser: vi.fn(),
+}));
+
 vi.mock("antd", async (importOriginal) => {
  const actual = await importOriginal<typeof import("antd")>();
  return {
@@ -71,10 +75,12 @@ vi.mock("antd", async (importOriginal) => {
 import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels";
 import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
 import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams";
+import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";

 const mockUseAllProxyModels = vi.mocked(useAllProxyModels);
 const mockUseTeam = vi.mocked(useTeam);
 const mockUseOrganization = vi.mocked(useOrganization);
+const mockUseCurrentUser = vi.mocked(useCurrentUser);

 describe("ModelSelect", () => {
  const mockProxyModels: ProxyModel[] = [
@@ -100,10 +106,16 @@ describe("ModelSelect", () => {
      data: undefined,
      isLoading: false,
    } as any);
+    mockUseCurrentUser.mockReturnValue({
+      data: { models: [] },
+      isLoading: false,
+    } as any);
  });

  it("should render", async () => {
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByTestId("model-select")).toBeInTheDocument();
@@ -116,7 +128,7 @@ describe("ModelSelect", () => {
      isLoading: true,
    } as any);

-    renderWithProviders(<ModelSelect onChange={mockOnChange} />);
+    renderWithProviders(<ModelSelect onChange={mockOnChange} context="user" />);

    expect(screen.getByTestId("skeleton-input")).toBeInTheDocument();
    expect(screen.queryByTestId("model-select")).not.toBeInTheDocument();
@@ -128,7 +140,7 @@ describe("ModelSelect", () => {
      isLoading: true,
    } as any);

-    renderWithProviders(<ModelSelect onChange={mockOnChange} teamID="team-1" />);
+    renderWithProviders(<ModelSelect onChange={mockOnChange} context="team" teamID="team-1" />);

    expect(screen.getByTestId("skeleton-input")).toBeInTheDocument();
  });
@@ -139,13 +151,49 @@ describe("ModelSelect", () => {
      isLoading: true,
    } as any);

-    renderWithProviders(<ModelSelect onChange={mockOnChange} organizationID="org-1" />);
+    renderWithProviders(<ModelSelect onChange={mockOnChange} context="organization" organizationID="org-1" />);
+
+    expect(screen.getByTestId("skeleton-input")).toBeInTheDocument();
+  });
+
+  it("should show skeleton loader when current user is loading", () => {
+    mockUseCurrentUser.mockReturnValue({
+      data: undefined,
+      isLoading: true,
+    } as any);
+
+    renderWithProviders(<ModelSelect onChange={mockOnChange} context="user" />);

    expect(screen.getByTestId("skeleton-input")).toBeInTheDocument();
  });

  it("should render special options group", async () => {
-    renderWithProviders(<ModelSelect onChange={mockOnChange} />);
+    const mockOrganization: Organization = {
+      organization_id: "org-1",
+      organization_alias: "Test Org",
+      budget_id: "budget-1",
+      metadata: {},
+      models: ["all-proxy-models"],
+      spend: 0,
+      model_spend: {},
+      created_at: "2024-01-01",
+      created_by: "user-1",
+      updated_at: "2024-01-01",
+      updated_by: "user-1",
+      litellm_budget_table: null,
+      teams: null,
+      users: null,
+      members: null,
+    };
+
+    mockUseOrganization.mockReturnValue({
+      data: mockOrganization,
+      isLoading: false,
+    } as any);
+
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="organization" organizationID="org-1" options={{ includeSpecialOptions: true }} />,
+    );

    await waitFor(() => {
      const select = screen.getByTestId("model-select");
@@ -156,7 +204,9 @@ describe("ModelSelect", () => {
  });

  it("should render wildcard options group", async () => {
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByText("All Openai models")).toBeInTheDocument();
@@ -165,7 +215,9 @@ describe("ModelSelect", () => {
  });

  it("should render regular models group", async () => {
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByText("gpt-4")).toBeInTheDocument();
@@ -175,7 +227,9 @@ describe("ModelSelect", () => {

  it("should call onChange when selecting a regular model", async () => {
    const user = userEvent.setup();
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByTestId("model-select")).toBeInTheDocument();
@@ -189,7 +243,37 @@ describe("ModelSelect", () => {

  it("should call onChange with only last special option when multiple special options are selected", async () => {
    const user = userEvent.setup();
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    const mockOrganization: Organization = {
+      organization_id: "org-1",
+      organization_alias: "Test Org",
+      budget_id: "budget-1",
+      metadata: {},
+      models: ["all-proxy-models"],
+      spend: 0,
+      model_spend: {},
+      created_at: "2024-01-01",
+      created_by: "user-1",
+      updated_at: "2024-01-01",
+      updated_by: "user-1",
+      litellm_budget_table: null,
+      teams: null,
+      users: null,
+      members: null,
+    };
+
+    mockUseOrganization.mockReturnValue({
+      data: mockOrganization,
+      isLoading: false,
+    } as any);
+
+    renderWithProviders(
+      <ModelSelect
+        onChange={mockOnChange}
+        context="organization"
+        organizationID="org-1"
+        options={{ showAllProxyModelsOverride: true, includeSpecialOptions: true }}
+      />,
+    );

    await waitFor(() => {
      expect(screen.getByTestId("model-select")).toBeInTheDocument();
@@ -203,7 +287,12 @@ describe("ModelSelect", () => {

  it("should disable regular models when special option is selected", async () => {
    renderWithProviders(
-      <ModelSelect onChange={mockOnChange} value={["all-proxy-models"]} showAllProxyModelsOverride={true} />,
+      <ModelSelect
+        onChange={mockOnChange}
+        value={["all-proxy-models"]}
+        context="user"
+        options={{ showAllProxyModelsOverride: true }}
+      />,
    );

    await waitFor(() => {
@@ -214,7 +303,12 @@ describe("ModelSelect", () => {

  it("should disable wildcard models when special option is selected", async () => {
    renderWithProviders(
-      <ModelSelect onChange={mockOnChange} value={["all-proxy-models"]} showAllProxyModelsOverride={true} />,
+      <ModelSelect
+        onChange={mockOnChange}
+        value={["all-proxy-models"]}
+        context="user"
+        options={{ showAllProxyModelsOverride: true }}
+      />,
    );

    await waitFor(() => {
@@ -224,8 +318,37 @@ describe("ModelSelect", () => {
  });

  it("should disable other special options when one special option is selected", async () => {
+    const mockOrganization: Organization = {
+      organization_id: "org-1",
+      organization_alias: "Test Org",
+      budget_id: "budget-1",
+      metadata: {},
+      models: ["all-proxy-models"],
+      spend: 0,
+      model_spend: {},
+      created_at: "2024-01-01",
+      created_by: "user-1",
+      updated_at: "2024-01-01",
+      updated_by: "user-1",
+      litellm_budget_table: null,
+      teams: null,
+      users: null,
+      members: null,
+    };
+
+    mockUseOrganization.mockReturnValue({
+      data: mockOrganization,
+      isLoading: false,
+    } as any);
+
    renderWithProviders(
-      <ModelSelect onChange={mockOnChange} value={["all-proxy-models"]} showAllProxyModelsOverride={true} />,
+      <ModelSelect
+        onChange={mockOnChange}
+        value={["all-proxy-models"]}
+        context="organization"
+        organizationID="org-1"
+        options={{ showAllProxyModelsOverride: true, includeSpecialOptions: true }}
+      />,
    );

    await waitFor(() => {
@@ -235,7 +358,9 @@ describe("ModelSelect", () => {
  });

  it("should filter models when showAllProxyModelsOverride is true", async () => {
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByText("gpt-4")).toBeInTheDocument();
@@ -267,7 +392,9 @@ describe("ModelSelect", () => {
      isLoading: false,
    } as any);

-    renderWithProviders(<ModelSelect onChange={mockOnChange} organizationID="org-1" />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="organization" organizationID="org-1" />,
+    );

    await waitFor(() => {
      expect(screen.getByText("gpt-4")).toBeInTheDocument();
@@ -275,7 +402,7 @@ describe("ModelSelect", () => {
    });
  });

-  it("should return empty models array when organization does not have all-proxy-models", async () => {
+  it("should filter models when organization has specific models", async () => {
    const mockOrganization: Organization = {
      organization_id: "org-1",
      organization_alias: "Test Org",
@@ -299,17 +426,24 @@ describe("ModelSelect", () => {
      isLoading: false,
    } as any);

-    renderWithProviders(<ModelSelect onChange={mockOnChange} organizationID="org-1" />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="organization" organizationID="org-1" />,
+    );

    await waitFor(() => {
-      expect(screen.queryByText("gpt-4")).not.toBeInTheDocument();
+      expect(screen.getByText("gpt-4")).toBeInTheDocument();
      expect(screen.queryByText("claude-3")).not.toBeInTheDocument();
    });
  });

  it("should use custom dataTestId when provided", async () => {
    renderWithProviders(
-      <ModelSelect onChange={mockOnChange} dataTestId="custom-test-id" showAllProxyModelsOverride={true} />,
+      <ModelSelect
+        onChange={mockOnChange}
+        dataTestId="custom-test-id"
+        context="user"
+        options={{ showAllProxyModelsOverride: true }}
+      />,
    );

    await waitFor(() => {
@@ -319,7 +453,9 @@ describe("ModelSelect", () => {

  it("should handle multiple model selections", async () => {
    const user = userEvent.setup();
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByTestId("model-select")).toBeInTheDocument();
@@ -337,7 +473,9 @@ describe("ModelSelect", () => {
  });

  it("should capitalize provider name in wildcard options", async () => {
-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      expect(screen.getByText("All Openai models")).toBeInTheDocument();
@@ -356,11 +494,74 @@ describe("ModelSelect", () => {
      isLoading: false,
    } as any);

-    renderWithProviders(<ModelSelect onChange={mockOnChange} showAllProxyModelsOverride={true} />);
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ showAllProxyModelsOverride: true }} />,
+    );

    await waitFor(() => {
      const gpt4Options = screen.getAllByText("gpt-4");
      expect(gpt4Options.length).toBeGreaterThan(0);
    });
  });
+
+  it("should filter models based on user context with includeUserModels option", async () => {
+    mockUseCurrentUser.mockReturnValue({
+      data: { models: ["gpt-4"] },
+      isLoading: false,
+    } as any);
+
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="user" options={{ includeUserModels: true }} />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("gpt-4")).toBeInTheDocument();
+      expect(screen.queryByText("claude-3")).not.toBeInTheDocument();
+    });
+  });
+
+  it("should filter models based on team context", async () => {
+    const mockTeam = {
+      team_id: "team-1",
+      team_alias: "Test Team",
+      models: ["gpt-4"],
+    };
+
+    const mockOrganization: Organization = {
+      organization_id: "org-1",
+      organization_alias: "Test Org",
+      budget_id: "budget-1",
+      metadata: {},
+      models: ["gpt-4"],
+      spend: 0,
+      model_spend: {},
+      created_at: "2024-01-01",
+      created_by: "user-1",
+      updated_at: "2024-01-01",
+      updated_by: "user-1",
+      litellm_budget_table: null,
+      teams: null,
+      users: null,
+      members: null,
+    };
+
+    mockUseTeam.mockReturnValue({
+      data: mockTeam,
+      isLoading: false,
+    } as any);
+
+    mockUseOrganization.mockReturnValue({
+      data: mockOrganization,
+      isLoading: false,
+    } as any);
+
+    renderWithProviders(
+      <ModelSelect onChange={mockOnChange} context="team" teamID="team-1" organizationID="org-1" />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("gpt-4")).toBeInTheDocument();
+      expect(screen.queryByText("claude-3")).not.toBeInTheDocument();
+    });
+  });
 });
@@ -1,77 +1,110 @@
 import { ProxyModel, useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels";
+import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
 import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams";
+import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
 import { Select, Skeleton, Tooltip, type SelectProps } from "antd";
 import { Organization, Team } from "../networking";
-import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
 import { splitWildcardModels } from "./modelUtils";

-const MODEL_SELECT_SPECIAL_VALUES = {
-  ALL_PROXY_MODELS: {
-    label: "All Proxy Models",
-    value: "all-proxy-models",
-  },
-  NO_DEFAULT_MODELS: {
-    label: "No Default Models",
-    value: "no-default-models",
-  },
-};
+const MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE = {
+  label: "All Proxy Models",
+  value: "all-proxy-models",
+} as const;

-const MODEL_SELECT_SPECIAL_VALUES_ARRAY = Object.values(MODEL_SELECT_SPECIAL_VALUES);
+const MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE = {
+  label: "No Default Models",
+  value: "no-default-models",
+} as const;

-export interface ModelSelectContext {
+const MODEL_SELECT_SPECIAL_VALUES_ARRAY = [
+  MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE,
+  MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE,
+] as const;
+
+export interface ModelSelectProps {
  teamID?: string;
  organizationID?: string;
-  includeUserModels?: boolean;
-  showAllTeamModelsOption?: boolean;
-  showAllProxyModelsOverride?: boolean;
-  includeSpecialOptions?: boolean;
+  options?: {
+    includeUserModels?: boolean;
+    showAllTeamModelsOption?: boolean;
+    showAllProxyModelsOverride?: boolean;
+    includeSpecialOptions?: boolean;
+  };
+  context: "team" | "organization" | "user";
  dataTestId?: string;
  value?: string[];
  onChange: (values: string[]) => void;
 }

-const filterModels = (
-  allProxyModels: ProxyModel[],
-  ctx: ModelSelectContext,
-  {
-    selectedTeam,
-    selectedOrganization,
-    userModels,
-  }: { selectedTeam?: Team; selectedOrganization?: Organization; userModels?: ProxyModel[] },
-): ProxyModel[] => {
-  const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((model) => [model.id, model])).values());
-  if (ctx.showAllProxyModelsOverride) {
-    return deduplicatedProxyModels;
-  }
-
-  if (selectedOrganization) {
-    if (selectedOrganization.models.includes(MODEL_SELECT_SPECIAL_VALUES.ALL_PROXY_MODELS.value)) {
-      return deduplicatedProxyModels;
-    }
-  }
-
-  return [];
+type FilterContextArgs = {
+  allProxyModels: string[];
+  selectedTeam?: Team;
+  selectedOrganization?: Organization;
+  userModels?: string[];
+  options?: ModelSelectProps["options"];
 };

-export const ModelSelect = (ctx: ModelSelectContext) => {
-  const {
-    teamID,
-    organizationID,
-    includeUserModels,
-    showAllTeamModelsOption,
-    showAllProxyModelsOverride,
-    includeSpecialOptions,
-    dataTestId,
-    value = [],
-    onChange,
-  } = ctx;
+const contextFilters: Record<ModelSelectProps["context"], (args: FilterContextArgs) => string[]> = {
+  user: ({ allProxyModels, userModels, options }) => {
+    if (!userModels) return [];
+    if (options?.includeUserModels) return userModels;
+    return [];
+  },
+
+  team: ({ allProxyModels, selectedOrganization, userModels }) => {
+    if (selectedOrganization) {
+      if (selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value)) {
+        return allProxyModels;
+      }
+      // Return organization's models (filtered from allProxyModels)
+      return allProxyModels.filter((model) => selectedOrganization.models.includes(model));
+    }
+
+    return userModels ?? [];
+  },
+
+  organization: ({ allProxyModels, selectedOrganization, options }) => {
+    if (!selectedOrganization) return [];
+
+    if (selectedOrganization.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value)) {
+      return allProxyModels;
+    }
+
+    return allProxyModels.filter((model) => selectedOrganization.models.includes(model));
+  },
+};
+
+const filterModels = (
+  allProxyModels: ProxyModel[],
+  ctx: ModelSelectProps,
+  extra: { selectedTeam?: Team; selectedOrganization?: Organization; userModels?: string[] },
+): string[] => {
+  const deduplicatedProxyModels = Array.from(new Map(allProxyModels.map((m) => [m.id, m])).values()).map(
+    (model) => model.id,
+  );
+  if (ctx.options?.showAllProxyModelsOverride) return deduplicatedProxyModels;
+
+  const filterFn = contextFilters[ctx.context];
+  if (!filterFn) return [];
+
+  return filterFn({ allProxyModels: deduplicatedProxyModels, ...extra, options: ctx.options });
+};
+
+export const ModelSelect = (props: ModelSelectProps) => {
+  const { teamID, organizationID, options, context, dataTestId, value = [], onChange } = props;
+  const { includeUserModels, showAllTeamModelsOption, showAllProxyModelsOverride, includeSpecialOptions } =
+    options || {};
  const { data: allProxyModels, isLoading: isLoadingAllProxyModels } = useAllProxyModels();
  const { data: team, isLoading: isLoadingTeam } = useTeam(teamID);
  const { data: organization, isLoading: isLoadingOrganization } = useOrganization(organizationID);
+  const { data: currentUser, isLoading: isCurrentUserLoading } = useCurrentUser();

  const isSpecialOption = (value: string) => MODEL_SELECT_SPECIAL_VALUES_ARRAY.some((sv) => sv.value === value);
  const hasSpecialOptionSelected = value.some(isSpecialOption);
-  const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization;
+  const isLoading = isLoadingAllProxyModels || isLoadingTeam || isLoadingOrganization || isCurrentUserLoading;
+  const shouldShowAllProxyModels =
+    showAllProxyModelsOverride ||
+    (organization?.models.includes(MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value) && includeSpecialOptions);

  if (isLoading) {
    return <Skeleton.Input active block />;
@@ -95,9 +128,10 @@ export const ModelSelect = (ctx: ModelSelectContext) => {
    onChange(finalValues);
  };

-  const filteredModels = filterModels(allProxyModels?.data ?? [], ctx, {
+  const filteredModels = filterModels(allProxyModels?.data ?? [], props, {
    selectedTeam: team,
    selectedOrganization: organization,
+    userModels: currentUser?.models,
  });

  const { wildcard, regular } = splitWildcardModels(filteredModels);
@@ -107,36 +141,60 @@ export const ModelSelect = (ctx: ModelSelectContext) => {
      value={value}
      onChange={handleChange}
      options={[
-        {
-          label: <span>Special Options</span>,
-          title: "Special Options",
-          options: MODEL_SELECT_SPECIAL_VALUES_ARRAY.map((specialValue) => ({
-            label: <span>{specialValue.label}</span>,
-            value: specialValue.value,
-            disabled: value.length > 0 && value.some((v) => isSpecialOption(v) && v !== specialValue.value),
-            key: specialValue.value,
-          })),
-        },
-        {
-          label: <span>Wildcard Options</span>,
-          title: "Wildcard Options",
-          options: wildcard.map((model) => {
-            const provider = model.id.replace("/*", "");
-            const capitalizedProvider = provider.charAt(0).toUpperCase() + provider.slice(1);
+        includeSpecialOptions
+          ? {
+              label: <span>Special Options</span>,
+              title: "Special Options",
+              options: [
+                ...(shouldShowAllProxyModels
+                  ? [
+                      {
+                        label: <span>All Proxy Models</span>,
+                        value: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
+                        disabled:
+                          value.length > 0 &&
+                          value.some(
+                            (v) => isSpecialOption(v) && v !== MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
+                          ),
+                        key: MODEL_SELECT_ALL_PROXY_MODELS_SPECIAL_VALUE.value,
+                      },
+                    ]
+                  : []),
+                {
+                  label: <span>No Default Models</span>,
+                  value: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value,
+                  disabled:
+                    value.length > 0 &&
+                    value.some((v) => isSpecialOption(v) && v !== MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value),
+                  key: MODEL_SELECT_NO_DEFAULT_MODELS_SPECIAL_VALUE.value,
+                },
+              ],
+            }
+          : [],
+        ...(wildcard.length > 0
+          ? [
+              {
+                label: <span>Wildcard Options</span>,
+                title: "Wildcard Options",
+                options: wildcard.map((model) => {
+                  const provider = model.replace("/*", "");
+                  const capitalizedProvider = provider.charAt(0).toUpperCase() + provider.slice(1);

-            return {
-              label: <span>{`All ${capitalizedProvider} models`}</span>,
-              value: model.id,
-              disabled: hasSpecialOptionSelected,
-            };
-          }),
-        },
+                  return {
+                    label: <span>{`All ${capitalizedProvider} models`}</span>,
+                    value: model,
+                    disabled: hasSpecialOptionSelected,
+                  };
+                }),
+              },
+            ]
+          : []),
        {
          label: <span>Models</span>,
          title: "Models",
          options: regular.map((model) => ({
-            label: <span>{model.id}</span>,
-            value: model.id,
+            label: <span>{model}</span>,
+            value: model,
            disabled: hasSpecialOptionSelected,
          })),
        },
@@ -1,4 +1,3 @@
-import type { ProxyModel } from "@/app/(dashboard)/hooks/models/useModels";
 import { describe, expect, it } from "vitest";
 import { splitWildcardModels } from "./modelUtils";

@@ -9,29 +8,21 @@ describe("splitWildcardModels", () => {
  });

  it("should split models into wildcard and regular groups", () => {
-    const models: ProxyModel[] = [
-      { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" },
-      { id: "openai/*", object: "model", created: 1234567890, owned_by: "openai" },
-      { id: "claude-3", object: "model", created: 1234567890, owned_by: "anthropic" },
-      { id: "anthropic/*", object: "model", created: 1234567890, owned_by: "anthropic" },
-    ];
+    const models: string[] = ["gpt-4", "openai/*", "claude-3", "anthropic/*"];

    const result = splitWildcardModels(models);

    expect(result.wildcard).toHaveLength(2);
-    expect(result.wildcard[0].id).toBe("openai/*");
-    expect(result.wildcard[1].id).toBe("anthropic/*");
+    expect(result.wildcard[0]).toBe("openai/*");
+    expect(result.wildcard[1]).toBe("anthropic/*");

    expect(result.regular).toHaveLength(2);
-    expect(result.regular[0].id).toBe("gpt-4");
-    expect(result.regular[1].id).toBe("claude-3");
+    expect(result.regular[0]).toBe("gpt-4");
+    expect(result.regular[1]).toBe("claude-3");
  });

  it("should return only wildcard models when all models are wildcard", () => {
-    const models: ProxyModel[] = [
-      { id: "openai/*", object: "model", created: 1234567890, owned_by: "openai" },
-      { id: "anthropic/*", object: "model", created: 1234567890, owned_by: "anthropic" },
-    ];
+    const models: string[] = ["openai/*", "anthropic/*"];

    const result = splitWildcardModels(models);

@@ -40,10 +31,7 @@ describe("splitWildcardModels", () => {
  });

  it("should return only regular models when no models are wildcard", () => {
-    const models: ProxyModel[] = [
-      { id: "gpt-4", object: "model", created: 1234567890, owned_by: "openai" },
-      { id: "claude-3", object: "model", created: 1234567890, owned_by: "anthropic" },
-    ];
+    const models: string[] = ["gpt-4", "claude-3"];

    const result = splitWildcardModels(models);

@@ -52,16 +40,12 @@ describe("splitWildcardModels", () => {
  });

  it("should correctly identify wildcard models ending with /*", () => {
-    const models: ProxyModel[] = [
-      { id: "provider/*", object: "model", created: 1234567890, owned_by: "provider" },
-      { id: "not-wildcard", object: "model", created: 1234567890, owned_by: "provider" },
-      { id: "also-not/*/wildcard", object: "model", created: 1234567890, owned_by: "provider" },
-    ];
+    const models: string[] = ["provider/*", "not-wildcard", "also-not/*/wildcard"];

    const result = splitWildcardModels(models);

    expect(result.wildcard).toHaveLength(1);
-    expect(result.wildcard[0].id).toBe("provider/*");
+    expect(result.wildcard[0]).toBe("provider/*");
    expect(result.regular).toHaveLength(2);
  });
 });
@@ -1,16 +1,14 @@
-import { ProxyModel } from "@/app/(dashboard)/hooks/models/useModels";
-
 export interface GroupedModels {
-  wildcard: ProxyModel[];
-  regular: ProxyModel[];
+  wildcard: string[];
+  regular: string[];
 }

-export const splitWildcardModels = (models: ProxyModel[]): GroupedModels => {
-  const wildcard: ProxyModel[] = [];
-  const regular: ProxyModel[] = [];
+export const splitWildcardModels = (models: string[]): GroupedModels => {
+  const wildcard: string[] = [];
+  const regular: string[] = [];

  for (const model of models) {
-    if (model.id.endsWith("/*")) {
+    if (model.endsWith("/*")) {
      wildcard.push(model);
    } else {
      regular.push(model);
@@ -29,12 +29,12 @@ import DeleteResourceModal from "./common_components/DeleteResourceModal";
 import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
 import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
 import MCPServerSelector from "./mcp_server_management/MCPServerSelector";
+import { ModelSelect } from "./ModelSelect/ModelSelect";
 import NotificationsManager from "./molecules/notifications_manager";
 import { Organization, organizationCreateCall, organizationDeleteCall, organizationListCall } from "./networking";
 import OrganizationInfoView from "./organization/organization_view";
 import NumericalInput from "./shared/numerical_input";
 import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
-import { ModelSelect } from "./ModelSelect/ModelSelect";

 interface OrganizationsTableProps {
  organizations: Organization[];
@@ -475,10 +475,10 @@ const OrganizationsTable: React.FC<OrganizationsTableProps> = ({
          </Form.Item>
          <Form.Item label="Models" name="models">
            <ModelSelect
-              showAllProxyModelsOverride
-              includeSpecialOptions
+              options={{ showAllProxyModelsOverride: true, includeSpecialOptions: true }}
              value={form.getFieldValue("models")}
              onChange={(values) => form.setFieldValue("models", values)}
+              context="organization"
            />
          </Form.Item>

@@ -1,7 +1,7 @@
 import * as networking from "@/components/networking";
 import { act, fireEvent, screen, waitFor } from "@testing-library/react";
 import { renderWithProviders } from "../../../tests/test-utils";
-import { afterEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import TeamInfoView from "./team_info";

 // Mock the networking module
@@ -17,7 +17,54 @@ vi.mock("@/components/networking", () => ({
  organizationInfoCall: vi.fn(),
 }));

+// Mock hooks used by ModelSelect
+vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({
+  useAllProxyModels: vi.fn(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  useTeam: vi.fn(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({
+  useOrganization: vi.fn(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({
+  useCurrentUser: vi.fn(),
+}));
+
+import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels";
+import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
+import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams";
+import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
+
+const mockUseAllProxyModels = vi.mocked(useAllProxyModels);
+const mockUseTeam = vi.mocked(useTeam);
+const mockUseOrganization = vi.mocked(useOrganization);
+const mockUseCurrentUser = vi.mocked(useCurrentUser);
+
 describe("TeamInfoView", () => {
+  beforeEach(() => {
+    // Set up default mock implementations
+    mockUseAllProxyModels.mockReturnValue({
+      data: { data: [] },
+      isLoading: false,
+    } as any);
+    mockUseTeam.mockReturnValue({
+      data: undefined,
+      isLoading: false,
+    } as any);
+    mockUseOrganization.mockReturnValue({
+      data: undefined,
+      isLoading: false,
+    } as any);
+    mockUseCurrentUser.mockReturnValue({
+      data: { models: [] },
+      isLoading: false,
+    } as any);
+  });
+
  afterEach(() => {
    vi.clearAllMocks();
  });
@@ -158,7 +205,7 @@ describe("TeamInfoView", () => {
    });

    await waitFor(() => {
-      expect(screen.getByLabelText("Models")).toBeInTheDocument();
+      expect(screen.getByTestId("models-select")).toBeInTheDocument();
    });

    const allProxyModelsOption = screen.queryByText("All Proxy Models");
@@ -170,6 +217,40 @@ describe("TeamInfoView", () => {
    const organizationModels = ["gpt-4", "claude-3-opus"];
    const userModels = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus", "claude-2"];

+    // Mock all proxy models - should include all user models
+    const allProxyModels = userModels.map((id) => ({
+      id,
+      object: "model",
+      created: 1234567890,
+      owned_by: "openai",
+    }));
+
+    mockUseAllProxyModels.mockReturnValue({
+      data: { data: allProxyModels },
+      isLoading: false,
+    } as any);
+
+    mockUseCurrentUser.mockReturnValue({
+      data: { models: userModels },
+      isLoading: false,
+    } as any);
+
+    const organizationData = {
+      organization_id: organizationId,
+      organization_name: "Test Organization",
+      spend: 0,
+      max_budget: null,
+      models: organizationModels,
+      tpm_limit: null,
+      rpm_limit: null,
+      members: null,
+    };
+
+    mockUseOrganization.mockReturnValue({
+      data: organizationData,
+      isLoading: false,
+    } as any);
+
    vi.mocked(networking.teamInfoCall).mockResolvedValue({
      team_id: "123",
      team_info: {
@@ -206,16 +287,7 @@ describe("TeamInfoView", () => {
      team_memberships: [],
    });

-    vi.mocked(networking.organizationInfoCall).mockResolvedValue({
-      organization_id: organizationId,
-      organization_name: "Test Organization",
-      spend: 0,
-      max_budget: null,
-      models: organizationModels,
-      tpm_limit: null,
-      rpm_limit: null,
-      members: null,
-    });
+    vi.mocked(networking.organizationInfoCall).mockResolvedValue(organizationData);

    vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] });
    vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
@@ -253,26 +325,39 @@ describe("TeamInfoView", () => {
    });

    await waitFor(() => {
-      expect(screen.getByLabelText("Models")).toBeInTheDocument();
+      expect(screen.getByTestId("models-select")).toBeInTheDocument();
    });

-    const modelsSelect = screen.getByLabelText("Models");
+    // Find the Ant Design Select selector element to open the dropdown
+    // The data-testid is on the Select component, we need to find the selector inside it
+    const modelsSelectElement = screen.getByTestId("models-select");
+    const selectSelector = modelsSelectElement.querySelector(".ant-select-selector");
+    expect(selectSelector).toBeTruthy();
+
+    // Open the dropdown by clicking on the selector
    act(() => {
-      fireEvent.mouseDown(modelsSelect);
+      fireEvent.mouseDown(selectSelector!);
    });

-    await waitFor(() => {
-      const dropdownOptions = screen.getAllByRole("option");
-      const optionTexts = dropdownOptions.map((option) => option.textContent);
+    // Wait for dropdown to open - Ant Design renders options in a portal
+    await waitFor(
+      () => {
+        const dropdownOptions = document.querySelectorAll(".ant-select-item-option");
+        expect(dropdownOptions.length).toBeGreaterThan(0);
+      },
+      { timeout: 5000 },
+    );

-      organizationModels.forEach((model) => {
-        expect(optionTexts).toContain(model);
-      });
+    const dropdownOptions = document.querySelectorAll(".ant-select-item-option");
+    const optionTexts = Array.from(dropdownOptions).map((option) => option.textContent?.trim() || "");

-      const modelsNotInOrganization = userModels.filter((m) => !organizationModels.includes(m));
-      modelsNotInOrganization.forEach((model) => {
-        expect(optionTexts).not.toContain(model);
-      });
+    organizationModels.forEach((model) => {
+      expect(optionTexts).toContain(model);
+    });
+
+    const modelsNotInOrganization = userModels.filter((m) => !organizationModels.includes(m));
+    modelsNotInOrganization.forEach((model) => {
+      expect(optionTexts).not.toContain(model);
    });
  }, 10000);

@@ -1,3 +1,4 @@
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import UserSearchModal from "@/components/common_components/user_search_modal";
 import {
  getGuardrailsList,
@@ -12,6 +13,7 @@ import {
 } from "@/components/networking";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils";
+import { isProxyAdminRole } from "@/utils/roles";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import {
@@ -34,11 +36,13 @@ import React, { useEffect, useMemo, useState } from "react";
 import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
 import AgentSelector from "../agent_management/AgentSelector";
 import DeleteResourceModal from "../common_components/DeleteResourceModal";
+import DurationSelect from "../common_components/DurationSelect";
 import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
-import { getModelDisplayName, unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key";
+import { unfurlWildcardModelsInList } from "../key_team_helpers/fetch_available_models_team_key";
 import LoggingSettingsView from "../logging_settings_view";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
 import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
+import { ModelSelect } from "../ModelSelect/ModelSelect";
 import NotificationsManager from "../molecules/notifications_manager";
 import { fetchMCPAccessGroups } from "../networking";
 import ObjectPermissionsView from "../object_permissions_view";
@@ -48,7 +52,6 @@ import EditLoggingSettings from "./EditLoggingSettings";
 import MemberModal from "./EditMembership";
 import MemberPermissions from "./member_permissions";
 import TeamMembersComponent from "./team_member_view";
-import DurationSelect from "../common_components/DurationSelect";

 export interface TeamMembership {
  user_id: string;
@@ -173,8 +176,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
  const [isDeleting, setIsDeleting] = useState(false);
  const [isTeamSaving, setIsTeamSaving] = useState(false);
  const [organization, setOrganization] = useState<Organization | null>(null);
-
-  console.log("userModels in team info", userModels);
+  const { userRole } = useAuthorized();

  const canEditTeam = is_team_admin || is_proxy_admin;

@@ -696,47 +698,19 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                    name="models"
                    rules={[{ required: true, message: "Please select at least one model" }]}
                  >
-                    <Select mode="multiple" placeholder="Select models">
-                      {(() => {
-                        let shouldShowAllProxyModels = false;
-
-                        if (organization) {
-                          // Team is in an organization
-                          if (organization.models.length === 0 || organization.models.includes("all-proxy-models")) {
-                            // Organization has empty array [] or "all-proxy-models"
-                            shouldShowAllProxyModels = true;
-                          }
-                          // Otherwise (organization has specific models), don't show "all-proxy-models"
-                        } else {
-                          // Team is not in an organization
-                          shouldShowAllProxyModels = is_proxy_admin || userModels.includes("all-proxy-models");
-                        }
-
-                        return shouldShowAllProxyModels ? (
-                          <Select.Option key="all-proxy-models" value="all-proxy-models">
-                            All Proxy Models
-                          </Select.Option>
-                        ) : null;
-                      })()}
-                      {(() => {
-                        // Show "no-default-models" option if:
-                        // 1. Team is not in an organization, OR
-                        // 2. Team is in an organization and organization's models include "no-default-models"
-                        const shouldShowNoDefaultModels =
-                          !organization || organization.models.includes("no-default-models");
-
-                        return shouldShowNoDefaultModels ? (
-                          <Select.Option key="no-default-models" value="no-default-models">
-                            No Default Models
-                          </Select.Option>
-                        ) : null;
-                      })()}
-                      {Array.from(new Set(modelsToPick)).map((model, idx) => (
-                        <Select.Option key={idx} value={model}>
-                          {getModelDisplayName(model)}
-                        </Select.Option>
-                      ))}
-                    </Select>
+                    <ModelSelect
+                      value={form.getFieldValue("models") || []}
+                      onChange={(values) => form.setFieldValue("models", values)}
+                      teamID={teamId}
+                      organizationID={teamData?.team_info?.organization_id || undefined}
+                      options={{
+                        includeSpecialOptions: true,
+                        includeUserModels: !teamData?.team_info?.organization_id,
+                        showAllProxyModelsOverride: isProxyAdminRole(userRole) && !teamData?.team_info?.organization_id,
+                      }}
+                      context="team"
+                      dataTestId="models-select"
+                    />
                  </Form.Item>

                  <Form.Item label="Max Budget (USD)" name="max_budget">
@@ -5,6 +5,7 @@ export interface UserInfo {
  user_role: string;
  spend: number;
  max_budget: number | null;
+  models: string[];
  key_count: number;
  created_at: string;
  updated_at: string;