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;