From b6b76323e66605fe30fc208a77f86650bc033402 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 16:08:39 -0700 Subject: [PATCH] fix(ui): abort in-flight LLM streams when their components unmount Five components (playground ChatUI, the public chat page, the prompt editor conversation panel, the usage AI chat panel, and the compliance batch runner) wire their AbortController only to a user-facing stop or close control. Navigating away mid-stream left the fetch and its ReadableStream reader running until the server finished, holding the response body and component closures alive and firing setState on unmounted components. Each now aborts its controller in an unmount cleanup. The conversation panel's controller moves from useState to a ref since nothing rendered from it; cancel no longer forces a redundant re-render --- .../components/chat_ui/ChatUI.test.tsx | 59 ++++++++++++++ .../playground/components/chat_ui/ChatUI.tsx | 2 + .../complianceUI/ComplianceUI.test.tsx | 46 +++++++++++ .../components/complianceUI/ComplianceUI.tsx | 6 ++ .../useConversation.test.tsx | 54 +++++++++++++ .../conversation_panel/useConversation.ts | 18 +++-- .../components/UsageAIChatPanel.test.tsx | 24 +++++- .../components/UsageAIChatPanel.tsx | 6 ++ .../src/app/chat/page.test.tsx | 78 +++++++++++++++++++ ui/litellm-dashboard/src/app/chat/page.tsx | 2 + 10 files changed, 288 insertions(+), 7 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/conversation_panel/useConversation.test.tsx create mode 100644 ui/litellm-dashboard/src/app/chat/page.test.tsx 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;