From e6fe329c3b7e23f43e0867767abf7f3cd68d07f3 Mon Sep 17 00:00:00 2001
From: nehaaprasaad <neh960112rxs@gmail.com>
Date: Thu, 7 May 2026 08:23:05 +0530
Subject: [PATCH] fix(ui): allow personal key creation when user has direct
 model access

---
 .../organisms/create_key_button.test.tsx      | 41 ++++++++++++++++++-
 .../organisms/create_key_button.tsx           |  6 ++-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx
index 3ad59cb3693..3574bd4bc12 100644
--- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx
+++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx
@@ -3,12 +3,13 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
 import CreateKey from "./create_key_button";
 
-const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall } = vi.hoisted(() => {
+const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall, mockModelAvailableCall } = vi.hoisted(() => {
   const formStateRef = { current: {} as Record<string, any> };
   const mockKeyCreateCall = vi.fn().mockResolvedValue({
     key: "test-api-key",
     soft_budget: null,
   });
+  const mockModelAvailableCall = vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] });
   const formMock = {
     setFieldsValue: vi.fn((values: Record<string, any>) => {
       Object.assign(formStateRef.current, values);
@@ -28,6 +29,7 @@ const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyC
     radioGroupValueRef,
     formStateRef,
     mockKeyCreateCall,
+    mockModelAvailableCall,
   };
 });
 
@@ -176,7 +178,7 @@ vi.mock("antd", () => {
 
 vi.mock("../networking", () => ({
   keyCreateCall: mockKeyCreateCall,
-  modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }),
+  modelAvailableCall: mockModelAvailableCall,
   getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }),
   getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }),
   getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }),
@@ -329,6 +331,7 @@ describe("CreateKey", () => {
       key: "test-api-key",
       soft_budget: null,
     });
+    mockModelAvailableCall.mockResolvedValue({ data: [{ id: "gpt-4" }] });
   });
 
   it("should render the CreateKey component", () => {
@@ -471,6 +474,40 @@ describe("CreateKey", () => {
     });
   });
 
+  it("should not force team selection when no-default-models appears with other models", async () => {
+    mockModelAvailableCall.mockResolvedValue({
+      data: [{ id: "no-default-models" }, { id: "gpt-4" }],
+    });
+
+    renderWithProviders(<CreateKey {...defaultProps} />);
+
+    act(() => {
+      fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+    });
+
+    await waitFor(() => {
+      expect(screen.queryByText(/please select a team to continue configuring your virtual key/i)).not.toBeInTheDocument();
+    });
+
+    expect(screen.getByText("Key Details")).toBeInTheDocument();
+  });
+
+  it("should force team selection when no-default-models is the only available model", async () => {
+    mockModelAvailableCall.mockResolvedValue({
+      data: [{ id: "no-default-models" }],
+    });
+
+    renderWithProviders(<CreateKey {...defaultProps} />);
+
+    act(() => {
+      fireEvent.click(screen.getByRole("button", { name: /create new key/i }));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText(/please select a team to continue configuring your virtual key/i)).toBeInTheDocument();
+    });
+  });
+
   describe("organization dropdown", () => {
     it("should render the organization dropdown when modal is open", async () => {
       renderWithProviders(<CreateKey {...defaultProps} />);
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 29050907503..7da65b0fcf5 100644
--- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
+++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
@@ -364,8 +364,10 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
     }
   }, [autoOpenCreate, prefillData, teams, hasPrefilled, form, userRole]);
 
-  // Check if team selection is required
-  const isTeamSelectionRequired = modelsToPick.includes("no-default-models");
+  // Require team selection only when the user has no personal/default model access.
+  const hasConcreteModelAccess = modelsToPick.some((model) => model !== "no-default-models");
+  const isTeamSelectionRequired =
+    !selectedCreateKeyTeam && modelsToPick.includes("no-default-models") && !hasConcreteModelAccess;
   const isFormDisabled = isTeamSelectionRequired && !selectedCreateKeyTeam;
 
   const handleCreate = async (formValues: Record<string, any>) => {