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 9da3e3a4a08..df75a01e1b7 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 @@ -2,12 +2,17 @@ import { act, fireEvent, render, screen, waitFor } from "@testing-library/react" import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; +import * as chatCompletionModule from "@/components/llm_calls/chat_completion"; // Mock the fetchAvailableModels function vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), })); +vi.mock("@/components/llm_calls/chat_completion", () => ({ + makeOpenAIChatCompletionRequest: vi.fn(), +})); + // Mock other networking functions that cause errors vi.mock("@/components/networking", () => ({ tagListCall: vi.fn().mockResolvedValue({ data: [] }), @@ -375,6 +380,60 @@ describe("ChatUI", () => { expect(customProxyInput).toHaveValue(testProxyUrl); }); + it("aborts the in-flight streaming request when the component unmounts", async () => { + vi.mocked(chatCompletionModule.makeOpenAIChatCompletionRequest).mockImplementation(() => new Promise(() => {})); + + const { unmount } = render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + const selectModelLabel = screen.getByText("Select Model"); + const modelSelectContainer = selectModelLabel.closest("div"); + const modelSelect = modelSelectContainer?.querySelector(".ant-select-selector"); + expect(modelSelect).toBeTruthy(); + + await act(async () => { + fireEvent.mouseDown(modelSelect!); + }); + + await waitFor(() => { + expect(screen.getAllByText("Model 1").length).toBeGreaterThan(0); + }); + + const model1Options = screen.getAllByText("Model 1"); + await act(async () => { + fireEvent.click(model1Options[model1Options.length - 1]); + }); + + const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)"); + await act(async () => { + fireEvent.change(messageInput, { target: { value: "hello" } }); + fireEvent.keyDown(messageInput, { key: "Enter" }); + }); + + await waitFor(() => { + expect(chatCompletionModule.makeOpenAIChatCompletionRequest).toHaveBeenCalled(); + }); + + const signal = vi.mocked(chatCompletionModule.makeOpenAIChatCompletionRequest).mock.calls[0][5]; + expect(signal).toBeInstanceOf(AbortSignal); + expect(signal!.aborted).toBe(false); + + unmount(); + + expect(signal!.aborted).toBe(true); + }); + it("should enable search functionality for MCP server selector", async () => { render( = ({ } }, [chatHistory]); + useEffect(() => () => abortControllerRef.current?.abort(), []); + const handleKeyDown = (event: React.KeyboardEvent) => { if (event.key === "Enter" && !event.shiftKey) { event.preventDefault(); // Prevent default to avoid newline diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.test.tsx new file mode 100644 index 00000000000..2aabe67de8d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.test.tsx @@ -0,0 +1,46 @@ +import { fireEvent, screen } from "@testing-library/react"; +import { beforeAll, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "@/../tests/test-utils"; +import { testPoliciesAndGuardrails } from "@/components/networking"; +import ComplianceUI from "./ComplianceUI"; + +vi.mock("@/components/networking", () => ({ + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + testPoliciesAndGuardrails: vi.fn(), +})); + +vi.mock("@/components/policies/PolicySelector", () => ({ + default: () => null, + getPolicyOptionEntries: () => [], +})); + +vi.mock("@/components/llm_calls/chat_completion", () => ({ + makeOpenAIChatCompletionRequest: vi.fn(), +})); + +beforeAll(() => { + Element.prototype.scrollIntoView = vi.fn(); +}); + +describe("ComplianceUI", () => { + it("aborts the in-flight batch run when the component unmounts", () => { + let capturedSignal: AbortSignal | undefined; + vi.mocked(testPoliciesAndGuardrails).mockImplementation((_accessToken, _body, signal) => { + capturedSignal = signal; + return new Promise(() => {}); + }); + + const { unmount } = renderWithProviders(); + + fireEvent.click(screen.getAllByText("All")[0]); + fireEvent.click(screen.getByRole("button", { name: /Simulate/ })); + + expect(testPoliciesAndGuardrails).toHaveBeenCalledTimes(1); + expect(capturedSignal).toBeDefined(); + expect(capturedSignal?.aborted).toBe(false); + + unmount(); + + expect(capturedSignal?.aborted).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index 39346105f2a..00987ddd2a7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -182,6 +182,12 @@ export default function ComplianceUI({ messagesEndRef.current?.scrollIntoView({ behavior: "smooth" }); }, [quickTestMessages]); + useEffect(() => { + return () => { + batchAbortControllerRef.current?.abort(); + }; + }, []); + const allFrameworks: ComplianceFramework[] = (() => { if (customPrompts.length === 0) return frameworks; const fwMap = new Map>(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.test.tsx new file mode 100644 index 00000000000..135e13e70c0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.test.tsx @@ -0,0 +1,54 @@ +import { act, renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { useConversation } from "./useConversation"; + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + fromBackend: vi.fn(), + info: vi.fn(), + success: vi.fn(), + }, +})); + +vi.mock("../utils", () => ({ + convertToDotPrompt: vi.fn().mockReturnValue(""), + extractVariables: vi.fn().mockReturnValue([]), +})); + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: () => "http://localhost:4000", + getGlobalLitellmHeaderName: () => "Authorization", +})); + +afterEach(() => { + vi.unstubAllGlobals(); +}); + +describe("useConversation", () => { + it("aborts the in-flight request when the hook unmounts", () => { + let capturedSignal: AbortSignal | null | undefined; + vi.stubGlobal( + "fetch", + vi.fn((_input: RequestInfo, init?: RequestInit) => { + capturedSignal = init?.signal; + return new Promise(() => {}); + }), + ); + + const { result, unmount } = renderHook(() => useConversation({}, "sk-test")); + + act(() => { + result.current.setInputMessage("hello"); + }); + act(() => { + void result.current.handleSendMessage(); + }); + + expect(capturedSignal).toBeDefined(); + expect(capturedSignal?.aborted).toBe(false); + + unmount(); + + expect(capturedSignal?.aborted).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.ts index d7852ca8bb4..ad43aba5f96 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.ts @@ -11,9 +11,15 @@ export const useConversation = (prompt: any, accessToken: string | null) => { const [inputMessage, setInputMessage] = useState(""); const [variables, setVariables] = useState>({}); const [variablesFilled, setVariablesFilled] = useState(false); - const [abortController, setAbortController] = useState(null); + const abortControllerRef = useRef(null); const messagesEndRef = useRef(null); + useEffect(() => { + return () => { + abortControllerRef.current?.abort(); + }; + }, []); + const extractedVariables = extractVariables(prompt); const allVariablesFilled = extractedVariables.every( @@ -59,7 +65,7 @@ export const useConversation = (prompt: any, accessToken: string | null) => { setInputMessage(""); const controller = new AbortController(); - setAbortController(controller); + abortControllerRef.current = controller; setIsLoading(true); const startTime = Date.now(); @@ -189,14 +195,14 @@ export const useConversation = (prompt: any, accessToken: string | null) => { } } finally { setIsLoading(false); - setAbortController(null); + abortControllerRef.current = null; } }; const handleCancelRequest = () => { - if (abortController) { - abortController.abort(); - setAbortController(null); + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + abortControllerRef.current = null; setIsLoading(false); NotificationsManager.info("Request cancelled"); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx index 54bf4fc25ce..3fbaf660b56 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx @@ -1,6 +1,7 @@ -import { screen } from "@testing-library/react"; +import { fireEvent, screen } from "@testing-library/react"; import { beforeAll, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "@/../tests/test-utils"; +import { usageAiChatStream } from "@/components/networking"; import UsageAIChatPanel from "./UsageAIChatPanel"; beforeAll(() => { @@ -77,4 +78,25 @@ describe("UsageAIChatPanel", () => { expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full"); expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0"); }); + + it("should abort the in-flight stream when the panel unmounts", () => { + let capturedSignal: AbortSignal | undefined; + vi.mocked(usageAiChatStream).mockImplementation((...args: Parameters) => { + capturedSignal = args[8]; + return new Promise(() => {}); + }); + + const { unmount } = renderWithProviders(); + + fireEvent.change(screen.getByPlaceholderText("Ask about your usage..."), { target: { value: "hello" } }); + fireEvent.click(screen.getByText("Send")); + + expect(usageAiChatStream).toHaveBeenCalledTimes(1); + expect(capturedSignal).toBeDefined(); + expect(capturedSignal?.aborted).toBe(false); + + unmount(); + + expect(capturedSignal?.aborted).toBe(true); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx index 9d6b09353a2..325424b900c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx @@ -115,6 +115,12 @@ const UsageAIChatPanel: React.FC = ({ open, onClose, acce } }, [open]); + useEffect(() => { + return () => { + abortControllerRef.current?.abort(); + }; + }, []); + useEffect(() => { if (typeof messagesEndRef.current?.scrollIntoView === "function") { messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); diff --git a/ui/litellm-dashboard/src/app/chat/page.test.tsx b/ui/litellm-dashboard/src/app/chat/page.test.tsx new file mode 100644 index 00000000000..6aa5e90c57e --- /dev/null +++ b/ui/litellm-dashboard/src/app/chat/page.test.tsx @@ -0,0 +1,78 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import ChatConversationPage from "./page"; +import * as responsesApiModule from "@/components/llm_calls/responses_api"; + +const { mockUseChatShell } = vi.hoisted(() => ({ + mockUseChatShell: vi.fn(() => ({ + accessToken: "token-123", + userId: "user-1", + userEmail: "user@example.com", + selectedMCPServers: [] as string[], + setSelectedMCPServers: vi.fn(), + activeConversationId: "conv-1", + activeConversation: { id: "conv-1", messages: [] }, + storageUnavailable: false, + staleId: null, + createConversation: vi.fn(() => "conv-1"), + appendMessage: vi.fn(), + updateLastAssistantMessage: vi.fn(), + truncateFromMessage: vi.fn(), + })), +})); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ replace: vi.fn(), push: vi.fn() }), +})); +vi.mock("@/contexts/ChatShellContext", () => ({ useChatShell: mockUseChatShell })); +vi.mock("@/components/chat/ChatShell", () => ({ + getChatRoutes: () => ({ chats: "/chat", integrations: "/chat/integrations" }), +})); +vi.mock("@/components/chat/ChatMessages", () => ({ + default: () =>
, +})); +vi.mock("@/components/chat/MCPConnectPicker", () => ({ + default: () =>
, +})); +vi.mock("@/components/molecules/message_manager", () => ({ + default: { error: vi.fn(), success: vi.fn(), info: vi.fn() }, +})); +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "gpt-5.2" }]), +})); +vi.mock("@/components/llm_calls/responses_api", () => ({ + makeOpenAIResponsesRequest: vi.fn(), +})); + +describe("ChatConversationPage", () => { + beforeEach(() => { + vi.clearAllMocks(); + localStorage.clear(); + }); + + it("aborts the in-flight streaming request when the component unmounts", async () => { + vi.mocked(responsesApiModule.makeOpenAIResponsesRequest).mockImplementation(() => new Promise(() => {})); + + const { unmount } = render(); + + await waitFor(() => { + expect(screen.getByText("gpt-5.2")).toBeInTheDocument(); + }); + + const textarea = screen.getByPlaceholderText("How can I help you today?"); + fireEvent.change(textarea, { target: { value: "hello" } }); + fireEvent.keyDown(textarea, { key: "Enter" }); + + await waitFor(() => { + expect(responsesApiModule.makeOpenAIResponsesRequest).toHaveBeenCalled(); + }); + + const signal = vi.mocked(responsesApiModule.makeOpenAIResponsesRequest).mock.calls[0][5]; + expect(signal).toBeInstanceOf(AbortSignal); + expect(signal!.aborted).toBe(false); + + unmount(); + + expect(signal!.aborted).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx index b6dccef47c6..61e6be4fa54 100644 --- a/ui/litellm-dashboard/src/app/chat/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.tsx @@ -90,6 +90,8 @@ export default function ChatConversationPage() { if (staleId) router.replace(getChatRoutes().chats); }, [staleId, router]); + useEffect(() => () => abortControllerRef.current?.abort(), []); + // Load models useEffect(() => { if (!accessToken) return;