From 65acf9bfefe6d26acdff09d906a2926d2a044dd3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 16 Oct 2026 09:57:17 -0700 Subject: [PATCH] fix(ui): close the playground chat gaps found in self-review Resend only complete text exchanges, so empty replies, MCP direct calls and transcriptions no longer leak into chat history. A stopped request can no longer append output, attach its gateway trace to a newer reply or clear the newer request's loading state. MCP selections are only pruned once both the server and toolset lists load. The MCP tool picker lists toolsets again and calls the server that backs the chosen tool, and a saved model the current key cannot use is cleared when the model list loads, matching the old chat Also vendors the scroll-fade-x utility the attachment primitive uses --- .../chat/PlaygroundChat.integration.test.tsx | 161 +++++++++++++++++- .../components/chat/PlaygroundChat.tsx | 52 +++++- .../components/chat/dispatchRequest.test.ts | 12 +- .../components/chat/dispatchRequest.ts | 15 +- .../components/chat/usePlaygroundCatalog.ts | 3 +- .../components/chat/usePlaygroundSend.ts | 21 ++- ui/litellm-dashboard/src/app/globals.css | 61 ++++++- 7 files changed, 302 insertions(+), 23 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx index d5a5bfdc85a..4d5da9e97e9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.integration.test.tsx @@ -1,9 +1,11 @@ import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { beforeEach, describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders as render, testQueryClient } from "@/../tests/test-utils"; import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; +import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "@/components/networking"; +import { setSecureItem } from "@/utils/secureStorage"; import { PlaygroundChat } from "./PlaygroundChat"; vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ default: () => true })); @@ -37,6 +39,17 @@ const replyWith = args[1](text, args[2]); }; +const preload = (values: Readonly>) => + Object.entries(values).forEach(([key, value]) => sessionStorage.setItem(key, value)); + +const deferred = () => { + const control: { release: () => void } = { release: () => undefined }; + const promise = new Promise((resolve) => { + control.release = resolve; + }); + return { promise, release: () => control.release() }; +}; + const renderChat = () => render(); @@ -56,6 +69,10 @@ describe("PlaygroundChat", () => { ]); }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + it("sends the chosen model and guardrails with the session key", async () => { const user = userEvent.setup(); chatRequest.mockImplementation(replyWith("Hi there")); @@ -99,4 +116,146 @@ describe("PlaygroundChat", () => { { role: "user", content: "try again" }, ]); }); + + it("sends every saved request setting with the chat request", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(fetchMCPServers).mockResolvedValue([{ server_id: "srv-1", alias: "docs" }]); + preload({ + selectedModel: "gpt-fast", + selectedTags: JSON.stringify(["team-a"]), + selectedVectorStores: JSON.stringify(["vs-1"]), + selectedPolicies: JSON.stringify(["strict"]), + selectedMCPServers: JSON.stringify(["srv-1"]), + streamingEnabled: "false", + }); + setSecureItem("customHeaders", JSON.stringify([["x-team", "blue"]])); + renderChat(); + + await waitFor(() => expect(fetchMCPServers).toHaveBeenCalled()); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + const args = chatRequest.mock.calls[0] as ChatArgs; + expect(args[2]).toBe("gpt-fast"); + expect(args[4]).toEqual(["team-a"]); + expect(args[10]).toEqual(["vs-1"]); + expect(args[12]).toEqual(["strict"]); + expect(args[13]).toEqual(["srv-1"]); + expect(args[25]).toBe(false); + expect(args[26]).toMatchObject({ "x-team": "blue" }); + }); + + it("keeps the saved MCP selection when the server list cannot be fetched", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(fetchMCPServers).mockRejectedValue(new Error("forbidden")); + preload({ selectedModel: "gpt-fast", selectedMCPServers: JSON.stringify(["srv-1"]) }); + renderChat(); + + await waitFor(() => expect(fetchMCPServers).toHaveBeenCalled()); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[13]).toEqual(["srv-1"]); + }); + + it("drops a saved model that the current key cannot use", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + preload({ selectedModel: "retired-model" }); + renderChat(); + + await waitFor(() => expect(fetchAvailableModels).toHaveBeenCalled()); + await waitFor(() => expect(screen.getByRole("combobox", { name: "Model" })).toHaveTextContent("Select a model")); + await user.click(screen.getByRole("combobox", { name: "Model" })); + await user.click(await screen.findByRole("option", { name: "gpt-fast" })); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[2]).toBe("gpt-fast"); + }); + + it("calls a toolset tool on the server that backs it", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPToolsets).mockResolvedValue([ + { toolset_id: "ts-1", toolset_name: "Research", tools: [{ server_id: "srv-2", tool_name: "search" }] }, + ]); + vi.mocked(listMCPTools).mockResolvedValue({ + tools: [{ name: "search", inputSchema: { type: "object", properties: {} } }], + }); + vi.mocked(callMCPTool).mockResolvedValue({ content: [{ type: "text", text: "found it" }] }); + preload({ endpointType: "mcp", selectedMCPServers: JSON.stringify(["toolset:ts-1"]) }); + renderChat(); + + await user.click(await screen.findByRole("combobox", { name: "Tool" })); + await user.click(await screen.findByRole("option", { name: /search/ })); + await user.click(screen.getByRole("button", { name: "Send message" })); + + await waitFor(() => expect(callMCPTool).toHaveBeenCalledTimes(1)); + expect(listMCPTools).toHaveBeenCalledWith("session-token", "srv-2"); + expect(vi.mocked(callMCPTool).mock.calls[0]?.slice(1, 3)).toEqual(["srv-2", "search"]); + }); + + it("ignores a stopped request that finishes after a newer one started", async () => { + const user = userEvent.setup(); + const first = deferred(); + const second = deferred(); + chatRequest + .mockImplementationOnce(async (...args: ChatArgs) => { + await first.promise; + args[1]("late reply", args[2]); + }) + .mockImplementationOnce(async (...args: ChatArgs) => { + await second.promise; + args[1]("fresh reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + first.release(); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(2)); + expect(screen.getByRole("button", { name: "Stop" })).toBeInTheDocument(); + expect(screen.queryByText("late reply")).not.toBeInTheDocument(); + second.release(); + expect(await screen.findByText("fresh reply")).toBeInTheDocument(); + expect(screen.queryByText("late reply")).not.toBeInTheDocument(); + }); + + it("keeps a stopped request's gateway details off the newer reply", async () => { + const user = userEvent.setup(); + vi.stubGlobal("matchMedia", (query: string) => ({ + matches: true, + media: query, + addEventListener: () => undefined, + removeEventListener: () => undefined, + })); + const first = deferred(); + chatRequest + .mockImplementationOnce(async (...args: ChatArgs) => { + args[27]?.(new Headers({ "x-litellm-call-id": "call-stopped" })); + await first.promise; + }) + .mockImplementationOnce(async (...args: ChatArgs) => { + args[27]?.(new Headers({ "x-litellm-call-id": "call-fresh" })); + args[1]("fresh reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + expect(await screen.findByText("call-fresh")).toBeInTheDocument(); + first.release(); + + await waitFor(() => expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument()); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(screen.getByText("call-fresh")).toBeInTheDocument(); + expect(screen.queryByText("call-stopped")).not.toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx index 22a2e481549..fe721be53b2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/PlaygroundChat.tsx @@ -1,6 +1,6 @@ "use client"; -import { useMemo, useRef, useState } from "react"; +import { useRef, useState } from "react"; import { ChevronDown, PanelRight, RotateCcw, SquarePen } from "lucide-react"; import useCan from "@/app/(dashboard)/hooks/useCan"; import { generateCodeSnippet } from "@/components/chat_ui/CodeSnippets"; @@ -74,6 +74,14 @@ const reachableMcp = (servers: readonly MCPServer[], toolsets: readonly MCPTools ...toolsets.map((toolset) => `${TOOLSET_PREFIX}${toolset.toolset_id}`), ]); +const toolsetFor = (selection: string | null, toolsets: readonly MCPToolset[]): MCPToolset | null => + toolsets.find((toolset) => `${TOOLSET_PREFIX}${toolset.toolset_id}` === selection) ?? null; + +const backingServers = (selection: string | null, toolset: MCPToolset | null): readonly string[] => { + if (toolset) return [...new Set(toolset.tools.map((tool) => tool.server_id))]; + return selection ? [selection] : []; +}; + const prefixed = (options: readonly PickerOption[]): readonly PickerOption[] => options.map((option) => ({ ...option, value: `${POLICY_PREFIX}${option.value}` })); @@ -104,7 +112,7 @@ export function PlaygroundChat({ const canViewPolicies = useCan("viewPolicies"); const apiKey = (stored.keySource === "session" ? accessToken : stored.virtualKey.trim()) || null; const mcp = useMcpCatalog(apiKey); - const settings = withReachableMcp(stored, mcp.isLoading ? null : reachableMcp(mcp.servers, mcp.toolsets)); + const settings = withReachableMcp(stored, mcp.isReady ? reachableMcp(mcp.servers, mcp.toolsets) : null); const capabilities = capabilitiesFor(settings.endpoint); const chat = useChatHistory({ simplified: false }); const codeInterpreter = useCodeInterpreter(); @@ -114,17 +122,35 @@ export function PlaygroundChat({ const isWide = useIsWide(); const models = usePlaygroundModels(apiKey); + const [checkedModels, setCheckedModels] = useState(models.data); + if (models.data && models.data !== checkedModels) { + setCheckedModels(models.data); + const loaded = models.data; + update((current) => + current.model && !loaded.some((option) => option.model_group === current.model) + ? { ...current, model: null } + : current, + ); + } const agents = usePlaygroundAgents(apiKey, settings.proxyBaseUrl, settings.endpoint === EndpointType.A2A_AGENTS); const guardrails = useGuardrailOptions(accessToken); const policies = usePolicyOptions(accessToken, canViewPolicies); const vectorStores = useVectorStoreOptions(accessToken); const tags = useTagOptions(accessToken); - const directServer = settings.endpoint === EndpointType.MCP ? settings.mcpSelection[0] ?? null : null; - const directServerTools = useServerTools(apiKey, directServer ? [directServer] : []); - const directToolList = directServer ? directServerTools[directServer]?.tools ?? [] : []; + const directSelection = settings.endpoint === EndpointType.MCP ? settings.mcpSelection[0] ?? null : null; + const directToolset = toolsetFor(directSelection, mcp.toolsets); + const directServers = backingServers(directSelection, directToolset); + const directServerTools = useServerTools(apiKey, directServers); + const backingTools = directServers.flatMap((server) => directServerTools[server]?.tools ?? []); + const directToolList = directToolset + ? backingTools.filter((tool) => directToolset.tools.some((entry) => entry.tool_name === tool.name)) + : backingTools; const directTool = directToolList.find((tool) => tool.name === settings.directTool) ?? null; - const customHeaders = useMemo(() => customHeadersFromPairs(settings.headerPairs), [settings.headerPairs]); + const directServer = directToolset + ? directToolset.tools.find((entry) => entry.tool_name === settings.directTool)?.server_id ?? null + : directSelection; + const customHeaders = customHeadersFromPairs(settings.headerPairs); const sendDeps: SendDeps = { settings, @@ -247,13 +273,21 @@ export function PlaygroundChat({ label: "MCP servers", options: mcp.servers.map((server) => ({ value: server.server_id, label: serverName(server) })), }, + { + label: "Toolsets", + options: mcp.toolsets.map((toolset) => ({ + value: `${TOOLSET_PREFIX}${toolset.toolset_id}`, + label: toolset.toolset_name, + description: toolset.description, + })), + }, ]} - value={directServer} + value={directSelection} onValueChange={(server) => update({ mcpSelection: [server], directTool: null })} isLoading={mcp.isLoading} emptyText="No MCP servers available for this key" /> - {directServer && ( + {directSelection && ( update({ directTool: tool })} - isLoading={directServerTools[directServer]?.isLoading ?? false} + isLoading={directServers.some((server) => directServerTools[server]?.isLoading)} emptyText="This server has no tools" /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.test.ts index 06b15c2edfd..df90daeb989 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.test.ts @@ -19,6 +19,10 @@ describe("textHistory", () => { user("question that was cancelled"), user("draw a cat"), assistant("data:image/png;base64,AAA", { isImage: true }), + { role: "user", content: "", mcpToolCall: { tool: "search", arguments: {} } }, + assistant("tool output"), + user("think hard"), + assistant(""), user("second question"), assistant("second answer"), ]; @@ -26,12 +30,18 @@ describe("textHistory", () => { expect(textHistory(history)).toEqual([ { role: "user", content: "first question" }, { role: "assistant", content: "first answer" }, - { role: "user", content: "draw a cat" }, { role: "user", content: "second question" }, { role: "assistant", content: "second answer" }, ]); }); + it("drops a transcription turn that has no typed text", () => { + expect(textHistory([user(""), assistant("transcribed words"), user("hi"), assistant("hello")])).toEqual([ + { role: "user", content: "hi" }, + { role: "assistant", content: "hello" }, + ]); + }); + it("drops a trailing user turn that never got a reply", () => { expect(textHistory([user("hi"), assistant("hello"), user("still waiting")])).toEqual([ { role: "user", content: "hi" }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.ts index 318432a98ef..cac0303b614 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/dispatchRequest.ts @@ -47,15 +47,18 @@ const NON_TEXT_FLAGS = ["isImage", "isAudio", "isError", "mcpToolCall"] as const const isText = (message: MessageType): boolean => NON_TEXT_FLAGS.every((flag) => !message[flag]); -const isReply = (message: MessageType | undefined): boolean => message?.role === "assistant" && !message.isError; +const isFilled = (content: MessageType["content"]): boolean => typeof content !== "string" || content.trim() !== ""; -const unanswered = (history: readonly MessageType[], index: number): boolean => - history[index]?.role === "user" && !isReply(history[index + 1]); +const isTurn = (message: MessageType | undefined, role: "user" | "assistant"): boolean => + message?.role === role && isText(message) && isFilled(message.content); + +const inExchange = (history: readonly MessageType[], index: number): boolean => + isTurn(history[index], "user") + ? isTurn(history[index + 1], "assistant") + : isTurn(history[index], "assistant") && isTurn(history[index - 1], "user"); export const textHistory = (history: readonly MessageType[]): ApiMessage[] => - history - .filter((message, index) => isText(message) && !unanswered(history, index)) - .map(({ role, content }) => ({ role, content })); + history.filter((_, index) => inExchange(history, index)).map(({ role, content }) => ({ role, content })); const toolResultText = (result: { content?: { type: string; text?: string }[] } | null): string => { if (result?.content && result.content.length > 0) { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundCatalog.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundCatalog.ts index a3c358e95c6..e1c8b8d9df9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundCatalog.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundCatalog.ts @@ -42,7 +42,7 @@ export function useMcpCatalog(requestKey: string | null) { const toolsets = useQuery({ queryKey: ["playground", "mcpToolsets", requestKey], queryFn: async () => { - const response: unknown = await fetchMCPToolsets(requestKey ?? "").catch(() => []); + const response: unknown = await fetchMCPToolsets(requestKey ?? ""); return Array.isArray(response) ? (response as MCPToolset[]) : []; }, enabled: Boolean(requestKey), @@ -51,6 +51,7 @@ export function useMcpCatalog(requestKey: string | null) { servers: servers.data ?? EMPTY, toolsets: toolsets.data ?? EMPTY, isLoading: servers.isLoading || toolsets.isLoading, + isReady: servers.isSuccess && toolsets.isSuccess, refetchServers: servers.refetch, }; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundSend.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundSend.ts index 105c0c23563..caee01afa7f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundSend.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat/usePlaygroundSend.ts @@ -84,6 +84,16 @@ const withTrace = (history: MessageType[], from: number, trace: ReplyTrace): Mes return history.map((message, index) => (index === replyIndex ? { ...message, trace } : message)); }; +const untilAborted = (target: T, signal: AbortSignal): T => + new Proxy(target, { + get: (object, key, receiver) => { + const value: unknown = Reflect.get(object, key, receiver); + return typeof value === "function" + ? (...args: unknown[]) => (signal.aborted ? undefined : value(...args)) + : value; + }, + }); + export function usePlaygroundSend(deps: SendDeps) { const [text, setText] = useState(""); const [files, setFiles] = useState([]); @@ -193,8 +203,8 @@ export function usePlaygroundSend(deps: SendDeps) { mcpServers: deps.mcpServers, mcpToolsets: deps.mcpToolsets, directCall: direct ? { ...direct, arguments: toolArguments ?? {} } : null, - chat, - codeInterpreter: deps.codeInterpreter, + chat: untilAborted(chat, controller.signal), + codeInterpreter: untilAborted(deps.codeInterpreter, controller.signal), onHeaders: (headers) => { captured.headers = headers; }, @@ -204,6 +214,7 @@ export function usePlaygroundSend(deps: SendDeps) { }, }; await dispatchRequest(context); + if (controller.signal.aborted) return; chat.setChatHistory((current) => withTrace(current, history.length, trace())); } catch (error) { if (controller.signal.aborted) return; @@ -212,8 +223,10 @@ export function usePlaygroundSend(deps: SendDeps) { { role: "assistant", content: errorText(error), isError: true, model: requestedModel, trace: trace() }, ]); } finally { - if (abortRef.current === controller) abortRef.current = null; - setIsLoading(false); + if (abortRef.current === controller) { + abortRef.current = null; + setIsLoading(false); + } } }; diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index f4fff9859f0..1534badbc7e 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -533,7 +533,66 @@ } } -/* shimmer */ +@property --scroll-fade-s { + syntax: ""; + inherits: false; + initial-value: 0px; +} + +@theme inline { + @keyframes scroll-fade-reveal-s { + from { + --scroll-fade-s: 0px; + } + to { + --scroll-fade-s: var(--_scroll-fade-size-s, var(--scroll-fade-size, min(12%, calc(var(--spacing) * 10)))); + } + } +} + +@utility scroll-fade-x { + --_scroll-fade-size-s: var(--scroll-fade-s-size, var(--scroll-fade-size, min(12%, calc(var(--spacing) * 10)))); + --_scroll-fade-size-e: var(--scroll-fade-e-size, var(--scroll-fade-size, min(12%, calc(var(--spacing) * 10)))); + --scroll-fade-inline: linear-gradient( + to right, + transparent 0, + #000 var(--scroll-fade-s, 0px), + #000 calc(100% - var(--scroll-fade-e, 0px)), + transparent 100% + ); + &:where([dir="rtl"], [dir="rtl"] *) { + --scroll-fade-inline: linear-gradient( + to left, + transparent 0, + #000 var(--scroll-fade-s, 0px), + #000 calc(100% - var(--scroll-fade-e, 0px)), + transparent 100% + ); + } + -webkit-mask-image: var(--scroll-fade-mask, var(--scroll-fade-inline)); + mask-image: var(--scroll-fade-mask, var(--scroll-fade-inline)); + -webkit-mask-composite: source-in; + mask-composite: intersect; + -webkit-mask-repeat: no-repeat; + mask-repeat: no-repeat; + + @supports (animation-timeline: scroll()) { + animation: + scroll-fade-reveal-s 1ms ease-in-out, + scroll-fade-reveal-e 1ms ease-in-out; + animation-timeline: scroll(self inline), scroll(self inline); + animation-range: + 0 var(--scroll-fade-reveal, calc(var(--spacing) * 24)), + calc(100% - var(--scroll-fade-reveal, calc(var(--spacing) * 24))) 100%; + animation-fill-mode: both; + } + + @supports not (animation-timeline: scroll()) { + --scroll-fade-s: var(--_scroll-fade-size-s); + --scroll-fade-e: var(--_scroll-fade-size-e); + } +} + @property --shimmer-angle { syntax: ""; inherits: true;