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 4d5da9e97e9..f58d64aa5d7 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,10 @@ -import { screen, waitFor } from "@testing-library/react"; +import { act, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; 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 { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "@/components/networking"; import { setSecureItem } from "@/utils/secureStorage"; import { PlaygroundChat } from "./PlaygroundChat"; @@ -14,6 +15,8 @@ vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi vi.mock("@/components/llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn() })); +vi.mock("@/components/llm_calls/responses_api", () => ({ makeOpenAIResponsesRequest: vi.fn() })); + vi.mock("@/components/networking", () => ({ tagListCall: vi.fn().mockResolvedValue({}), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), @@ -29,8 +32,12 @@ vi.mock("@/components/networking", () => ({ const chatRequest = vi.mocked(makeOpenAIChatCompletionRequest); +const responsesRequest = vi.mocked(makeOpenAIResponsesRequest); + type ChatArgs = Parameters; +type ResponsesArgs = Parameters; + const sentMessages = (call: number) => (chatRequest.mock.calls[call] as ChatArgs)[0]; const replyWith = @@ -50,6 +57,17 @@ const deferred = () => { return { promise, release: () => control.release() }; }; +const mcpTool = (name: string, required: readonly string[] = []) => ({ + name, + inputSchema: { + type: "object", + properties: Object.fromEntries(required.map((field) => [field, { type: "string" }])), + required: [...required], + }, +}); + +const settled = () => waitFor(() => expect(testQueryClient.isFetching()).toBe(0)); + const renderChat = () => render(); @@ -64,6 +82,12 @@ describe("PlaygroundChat", () => { sessionStorage.clear(); testQueryClient.clear(); Element.prototype.scrollIntoView = vi.fn(); + chatRequest.mockReset(); + responsesRequest.mockReset(); + vi.mocked(callMCPTool).mockReset(); + vi.mocked(fetchMCPServers).mockReset().mockResolvedValue([]); + vi.mocked(fetchMCPToolsets).mockReset().mockResolvedValue([]); + vi.mocked(listMCPTools).mockReset().mockResolvedValue({ tools: [] }); vi.mocked(fetchAvailableModels).mockResolvedValue([ { model_group: "gpt-fast", mode: "chat", providers: ["openai"] }, ]); @@ -154,6 +178,7 @@ describe("PlaygroundChat", () => { renderChat(); await waitFor(() => expect(fetchMCPServers).toHaveBeenCalled()); + await settled(); await send(user, "hello"); await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); @@ -194,9 +219,142 @@ describe("PlaygroundChat", () => { await waitFor(() => expect(callMCPTool).toHaveBeenCalledTimes(1)); expect(listMCPTools).toHaveBeenCalledWith("session-token", "srv-2"); + expect(listMCPTools).not.toHaveBeenCalledWith("session-token", "toolset:ts-1"); expect(vi.mocked(callMCPTool).mock.calls[0]?.slice(1, 3)).toEqual(["srv-2", "search"]); }); + it("offers toolsets as targets and lists only the toolset's own tools", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPToolsets).mockResolvedValue([ + { + toolset_id: "ts-1", + toolset_name: "Research", + tools: [ + { server_id: "srv-a", tool_name: "search" }, + { server_id: "srv-b", tool_name: "fetch" }, + { server_id: "srv-b", tool_name: "search" }, + ], + }, + ]); + vi.mocked(listMCPTools).mockImplementation(async (_key: string, server: string) => ({ + tools: + server === "srv-a" ? [mcpTool("search"), mcpTool("fetch", ["a_arg"])] : [mcpTool("search"), mcpTool("fetch")], + })); + vi.mocked(callMCPTool).mockResolvedValue({ content: [{ type: "text", text: "fetched" }] }); + preload({ endpointType: "mcp" }); + renderChat(); + + await user.click(await screen.findByRole("combobox", { name: "Server" })); + await user.click(await screen.findByRole("option", { name: /Research/ })); + await settled(); + await user.click(screen.getByRole("combobox", { name: "Tool" })); + await waitFor(() => + expect(screen.getAllByRole("option").map((option) => option.textContent)).toEqual(["search", "fetch"]), + ); + await user.click(screen.getByRole("option", { name: "fetch" })); + await user.click(screen.getByRole("button", { name: "Send message" })); + + await waitFor(() => expect(callMCPTool).toHaveBeenCalledTimes(1)); + expect(vi.mocked(callMCPTool).mock.calls[0]?.slice(1, 4)).toEqual(["srv-b", "fetch", {}]); + }); + + it("says the tools could not be loaded when the server's tool list fails", async () => { + const user = userEvent.setup(); + vi.mocked(fetchMCPServers).mockResolvedValue([{ server_id: "srv-1", alias: "docs" }]); + vi.mocked(listMCPTools).mockResolvedValue({ error: "server unreachable" }); + preload({ endpointType: "mcp", selectedMCPServers: JSON.stringify(["srv-1"]) }); + renderChat(); + + await settled(); + await user.click(screen.getByRole("combobox", { name: "Tool" })); + + expect(await screen.findByText("Unable to load tools for this server")).toBeInTheDocument(); + }); + + it("keeps the saved MCP selection when the toolset list cannot be fetched", async () => { + const user = userEvent.setup(); + chatRequest.mockImplementation(replyWith("Hi there")); + vi.mocked(fetchMCPServers).mockResolvedValue([{ server_id: "srv-1", alias: "docs" }]); + vi.mocked(fetchMCPToolsets).mockRejectedValue(new Error("database not connected")); + preload({ selectedModel: "gpt-fast", selectedMCPServers: JSON.stringify(["srv-1", "srv-gone"]) }); + renderChat(); + + await waitFor(() => expect(fetchMCPToolsets).toHaveBeenCalled()); + await settled(); + await send(user, "hello"); + + await waitFor(() => expect(chatRequest).toHaveBeenCalledTimes(1)); + expect((chatRequest.mock.calls[0] as ChatArgs)[13]).toEqual(["srv-1", "srv-gone"]); + }); + + it("drops a saved model that is missing from an already loaded model list", async () => { + const first = renderChat(); + await waitFor(() => expect(fetchAvailableModels).toHaveBeenCalled()); + await settled(); + first.unmount(); + sessionStorage.setItem("selectedModel", "retired-model"); + + renderChat(); + + await waitFor(() => expect(screen.getByRole("combobox", { name: "Model" })).toHaveTextContent("Select a model")); + expect(fetchAvailableModels).toHaveBeenCalledTimes(1); + }); + + it("drops a request that is still running when a new chat starts", async () => { + const user = userEvent.setup(); + const first = deferred(); + chatRequest.mockImplementationOnce(async (...args: ChatArgs) => { + await first.promise; + args[1]("late reply", args[2]); + }); + preload({ selectedModel: "gpt-fast" }); + renderChat(); + + await send(user, "first"); + expect(await screen.findByRole("button", { name: "Stop" })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "New chat" })); + await act(async () => { + first.release(); + await chatRequest.mock.results[0]?.value; + }); + + expect(chatRequest).toHaveBeenCalledTimes(1); + expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument(); + expect(screen.queryByText("late reply")).not.toBeInTheDocument(); + }); + + it("keeps a stopped request's response id out of the next request", async () => { + const user = userEvent.setup(); + const first = deferred(); + responsesRequest + .mockImplementationOnce(async (...args: ResponsesArgs) => { + await first.promise; + args[15]?.("resp-stopped"); + }) + .mockImplementationOnce(async (...args: ResponsesArgs) => { + args[15]?.("resp-fresh"); + args[1]("assistant", "fresh reply", args[2]); + }) + .mockImplementationOnce(async (...args: ResponsesArgs) => { + args[1]("assistant", "third reply", args[2]); + }); + preload({ endpointType: "responses", selectedModel: "gpt-fast", useApiSessionManagement: "true" }); + renderChat(); + + await send(user, "first"); + await user.click(await screen.findByRole("button", { name: "Stop" })); + await send(user, "second"); + expect(await screen.findByText("fresh reply")).toBeInTheDocument(); + await act(async () => { + first.release(); + await responsesRequest.mock.results[0]?.value; + }); + await send(user, "third"); + + await waitFor(() => expect(responsesRequest).toHaveBeenCalledTimes(3)); + expect((responsesRequest.mock.calls[2] as ResponsesArgs)[14]).toBe("resp-fresh"); + }); + it("ignores a stopped request that finishes after a newer one started", async () => { const user = userEvent.setup(); const first = deferred(); @@ -251,10 +409,12 @@ describe("PlaygroundChat", () => { await user.click(await screen.findByRole("button", { name: "Stop" })); await send(user, "second"); expect(await screen.findByText("call-fresh")).toBeInTheDocument(); - first.release(); + await act(async () => { + first.release(); + await chatRequest.mock.results[0]?.value; + }); - await waitFor(() => expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument()); - await new Promise((resolve) => setTimeout(resolve, 0)); + expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument(); 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 fe721be53b2..5374efa5518 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 @@ -9,7 +9,7 @@ import type { MessageType } from "@/components/chat_ui/types"; import type { ModelGroup } from "@/components/llm_calls/fetch_models"; import { customHeadersFromPairs } from "@/components/llm_calls/request_headers"; import type { MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; -import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import type { MCPServer, MCPTool, MCPToolset } from "@/components/mcp_tools/types"; import { Button } from "@/components/ui/button"; import { ButtonGroup } from "@/components/ui/button-group"; import { @@ -43,6 +43,7 @@ import { useTagOptions, useVectorStoreOptions, type PickerOption, + type ServerTools, } from "./usePlaygroundCatalog"; import { usePlaygroundSend, type SendDeps } from "./usePlaygroundSend"; import { usePlaygroundSettings } from "./usePlaygroundSettings"; @@ -79,7 +80,26 @@ const toolsetFor = (selection: string | null, toolsets: readonly MCPToolset[]): 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] : []; + return selection && !selection.startsWith(TOOLSET_PREFIX) ? [selection] : []; +}; + +type DirectTarget = Readonly<{ server: string; tool: MCPTool }>; + +const toolsetTargets = (toolset: MCPToolset, loaded: Readonly>): DirectTarget[] => + toolset.tools + .flatMap((entry) => { + const tool = loaded[entry.server_id]?.tools.find((candidate) => candidate.name === entry.tool_name); + return tool ? [{ server: entry.server_id, tool }] : []; + }) + .filter((target, index, targets) => targets.findIndex((other) => other.tool.name === target.tool.name) === index); + +const directTargets = ( + servers: readonly string[], + toolset: MCPToolset | null, + loaded: Readonly>, +): readonly DirectTarget[] => { + if (toolset) return toolsetTargets(toolset, loaded); + return servers.flatMap((server) => (loaded[server]?.tools ?? []).map((tool) => ({ server, tool }))); }; const prefixed = (options: readonly PickerOption[]): readonly PickerOption[] => @@ -122,7 +142,7 @@ export function PlaygroundChat({ const isWide = useIsWide(); const models = usePlaygroundModels(apiKey); - const [checkedModels, setCheckedModels] = useState(models.data); + const [checkedModels, setCheckedModels] = useState(undefined); if (models.data && models.data !== checkedModels) { setCheckedModels(models.data); const loaded = models.data; @@ -142,14 +162,10 @@ export function PlaygroundChat({ 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 directServer = directToolset - ? directToolset.tools.find((entry) => entry.tool_name === settings.directTool)?.server_id ?? null - : directSelection; + const targets = directTargets(directServers, directToolset, directServerTools); + const directTarget = targets.find((target) => target.tool.name === settings.directTool) ?? null; + const directTool = directTarget?.tool ?? null; + const toolsFailed = directServers.some((server) => directServerTools[server]?.failed); const customHeaders = customHeadersFromPairs(settings.headerPairs); const sendDeps: SendDeps = { @@ -160,8 +176,7 @@ export function PlaygroundChat({ customHeaders, mcpServers: mcp.servers, mcpToolsets: mcp.toolsets, - directCall: () => - directServer && settings.directTool ? { serverId: directServer, tool: settings.directTool } : null, + directCall: () => (directTarget ? { serverId: directTarget.server, tool: directTarget.tool.name } : null), toolArgsForm, chat, codeInterpreter, @@ -294,7 +309,7 @@ export function PlaygroundChat({ groups={[ { label: "Tools", - options: directToolList.map((tool) => ({ + options: targets.map(({ tool }) => ({ value: tool.name, label: tool.name, description: tool.description, @@ -304,7 +319,7 @@ export function PlaygroundChat({ value={settings.directTool} onValueChange={(tool) => update({ directTool: tool })} isLoading={directServers.some((server) => directServerTools[server]?.isLoading)} - emptyText="This server has no tools" + emptyText={toolsFailed ? "Unable to load tools for this server" : "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 df90daeb989..e568f709411 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 @@ -23,6 +23,10 @@ describe("textHistory", () => { assistant("tool output"), user("think hard"), assistant(""), + user("think again"), + assistant(" \n"), + user("embed this"), + assistant("[0.1, 0.2]", { isEmbeddings: true }), user("second question"), assistant("second answer"), ]; 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 cac0303b614..3d432df315a 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 @@ -43,7 +43,7 @@ export type DispatchContext = Readonly<{ const optional = (values: readonly string[]): string[] | undefined => (values.length > 0 ? [...values] : undefined); -const NON_TEXT_FLAGS = ["isImage", "isAudio", "isError", "mcpToolCall"] as const; +const NON_TEXT_FLAGS = ["isImage", "isAudio", "isEmbeddings", "isError", "mcpToolCall"] as const; const isText = (message: MessageType): boolean => NON_TEXT_FLAGS.every((flag) => !message[flag]); 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 e1c8b8d9df9..bdb6424699e 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 @@ -39,10 +39,11 @@ export function useMcpCatalog(requestKey: string | null) { queryFn: async () => serverList(await fetchMCPServers(requestKey ?? "")), enabled: Boolean(requestKey), }); - const toolsets = useQuery({ + const toolsets = useQuery({ queryKey: ["playground", "mcpToolsets", requestKey], queryFn: async () => { - const response: unknown = await fetchMCPToolsets(requestKey ?? ""); + const response: unknown = await fetchMCPToolsets(requestKey ?? "").catch(() => null); + if (response === null) return null; return Array.isArray(response) ? (response as MCPToolset[]) : []; }, enabled: Boolean(requestKey), @@ -51,7 +52,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, + isReady: servers.isSuccess && Boolean(toolsets.data), 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 caee01afa7f..6a704430fd2 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 @@ -113,7 +113,14 @@ export function usePlaygroundSend(deps: SendDeps) { setFiles(files.filter((_, i) => i !== index)); }; + const stop = () => { + abortRef.current?.abort(); + abortRef.current = null; + setIsLoading(false); + }; + const clearDraft = () => { + stop(); files.forEach((file) => file.previewUrl && URL.revokeObjectURL(file.previewUrl)); setFiles([]); setText(""); @@ -121,9 +128,7 @@ export function usePlaygroundSend(deps: SendDeps) { const cancel = () => { if (!abortRef.current) return; - abortRef.current.abort(); - abortRef.current = null; - setIsLoading(false); + stop(); toast.info("Request cancelled"); }; @@ -181,6 +186,7 @@ export function usePlaygroundSend(deps: SendDeps) { const controller = new AbortController(); abortRef.current = controller; + const live = untilAborted(chat, controller.signal); const captured: { headers?: Headers; responseId?: string } = {}; const trace = (): ReplyTrace => ({ ...gatewayTraceFromHeaders(captured.headers ?? new Headers()), @@ -203,14 +209,14 @@ export function usePlaygroundSend(deps: SendDeps) { mcpServers: deps.mcpServers, mcpToolsets: deps.mcpToolsets, directCall: direct ? { ...direct, arguments: toolArguments ?? {} } : null, - chat: untilAborted(chat, controller.signal), - codeInterpreter: untilAborted(deps.codeInterpreter, controller.signal), + chat: live, + codeInterpreter: deps.codeInterpreter, onHeaders: (headers) => { captured.headers = headers; }, onResponseId: (responseId) => { captured.responseId = responseId; - chat.handleResponseId(responseId); + live.handleResponseId(responseId); }, }; await dispatchRequest(context);