From f4fb4e0bea79908cea742e007e642b2034b6e3d6 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 13 Aug 2026 12:43:25 -0700
Subject: [PATCH] fix(ui): label the virtual key source and unstick a cancelled
 model load

The key source trigger rendered the stored value, so the playground showed
session and custom instead of Current UI Session and Virtual Key. Name the
selected option on the trigger.

Clearing the key while models were loading left the selector disabled for
good: the in-flight load skips its reset once cancelled, and the branch that
handles an empty key returned without clearing the loading flag, so nothing
put it back. Clear it on that path too.
---
 .../components/chat_ui/ChatUI.test.tsx        | 67 +++++++++++++++++++
 .../playground/components/chat_ui/ChatUI.tsx  |  3 +-
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx
index 8e242e6ec35..32f1013f386 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx
@@ -475,4 +475,71 @@ describe("ChatUI", () => {
 
     expect(await screen.findByText("clip.wav")).toBeInTheDocument();
   });
+
+  it("should name the virtual key source options instead of showing raw values", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <ChatUI
+        accessToken="1234567890"
+        token="1234567890"
+        userRole="user"
+        userID="1234567890"
+        disabledPersonalKeyCreation={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Test Key")).toBeInTheDocument();
+    });
+
+    const keySourceTrigger = screen.getByLabelText("Virtual Key Source");
+    expect(keySourceTrigger).toHaveTextContent("Current UI Session");
+    expect(keySourceTrigger).not.toHaveTextContent("session");
+
+    await user.click(keySourceTrigger);
+    await user.click(await screen.findByRole("option", { name: "Virtual Key" }));
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Virtual Key Source")).toHaveTextContent("Virtual Key");
+    });
+    expect(screen.getByLabelText("Virtual Key Source")).not.toHaveTextContent("custom");
+  });
+
+  it("should re-enable the model selector when the virtual key is cleared mid-load", async () => {
+    const user = userEvent.setup();
+    (fetchModelsModule.fetchAvailableModels as ReturnType<typeof vi.fn>).mockImplementation(
+      () => new Promise(() => {}),
+    );
+
+    render(
+      <ChatUI
+        accessToken="1234567890"
+        token="1234567890"
+        userRole="user"
+        userID="1234567890"
+        disabledPersonalKeyCreation={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Test Key")).toBeInTheDocument();
+    });
+
+    await user.click(screen.getByLabelText("Virtual Key Source"));
+    await user.click(await screen.findByRole("option", { name: "Virtual Key" }));
+
+    const keyField = await screen.findByPlaceholderText("Enter custom Virtual Key");
+    await user.type(keyField, "sk-test");
+
+    await waitFor(() => {
+      expect(screen.getByPlaceholderText("Loading models...")).toBeInTheDocument();
+    });
+
+    await user.clear(keyField);
+
+    await waitFor(() => {
+      expect(screen.getByPlaceholderText("Select a Model")).not.toBeDisabled();
+    });
+  });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx
index f3d79dd3a2f..9ef44589240 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx
@@ -407,6 +407,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
     if (!userApiKey) {
       setModelInfo([]);
       setModelLoadError(false);
+      setIsLoadingModels(false);
       return;
     }
 
@@ -1240,7 +1241,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
                     }}
                   >
                     <SelectTrigger className="w-full" size="sm" aria-label="Virtual Key Source">
-                      <SelectValue />
+                      <SelectValue>{apiKeySource === "custom" ? "Virtual Key" : "Current UI Session"}</SelectValue>
                     </SelectTrigger>
                     <SelectContent>
                       <SelectItem value="session">Current UI Session</SelectItem>