import { describe, it, expect, beforeEach, vi } from "vitest"; import { render, screen, fireEvent, act } from "@testing-library/react"; import { MantineProvider } from "@mantine/core"; // Shared, hoisted mock state so the @ai-sdk/react and "ai" module mocks (hoisted // above the imports) can expose the captured useChat callbacks / transport and // the spies back to the test body. const h = vi.hoisted(() => ({ state: { status: "streaming" as string, onFinish: null as null | ((arg: Record) => void), sendMessage: vi.fn(), stop: vi.fn(), transport: null as null | { prepareSendMessagesRequest: (arg: { messages: unknown[]; body: Record; }) => { body: Record }; }, }, })); // Mock useChat: capture onFinish, return the spies and the controllable status. vi.mock("@ai-sdk/react", () => ({ useChat: (opts: { onFinish?: (arg: Record) => void }) => { h.state.onFinish = opts.onFinish ?? null; return { messages: [], sendMessage: h.state.sendMessage, status: h.state.status, stop: h.state.stop, error: null, }; }, })); // Mock "ai": deterministic ids + a transport that records its options so the test // can invoke prepareSendMessagesRequest and assert the `interrupted` flag. vi.mock("ai", () => { let counter = 0; return { generateId: () => `gid-${counter++}`, DefaultChatTransport: class { constructor(opts: { prepareSendMessagesRequest: (arg: { messages: unknown[]; body: Record; }) => { body: Record }; }) { h.state.transport = opts; } }, }; }); // Stub the heavy children: MessageList (markdown/render) and ChatInput (the // composer). The ChatInput stub exposes a button that queues a message, the only // interaction this test needs to populate the queue while "streaming". vi.mock("@/features/ai-chat/components/message-list.tsx", () => ({ default: () =>
, })); vi.mock("@/features/ai-chat/components/chat-input.tsx", () => ({ default: ({ onQueue }: { onQueue: (text: string) => void }) => ( ), })); import ChatThread from "./chat-thread"; function renderThread() { const onTurnFinished = vi.fn(); render( , ); return { onTurnFinished }; } describe("ChatThread — send now (#198)", () => { beforeEach(() => { h.state.status = "streaming"; h.state.onFinish = null; h.state.sendMessage.mockClear(); h.state.stop.mockClear(); h.state.transport = null; }); it("aborts the current turn and resends the queued message on the abort", () => { renderThread(); // Queue a message while the turn is streaming. fireEvent.click(screen.getByTestId("queue-btn")); const sendNowBtn = screen.getByLabelText("Send now"); expect(sendNowBtn).toBeTruthy(); // "Send now" interrupts the current turn (stop), but does NOT send yet — // the resend happens once the abort lands in onFinish. fireEvent.click(sendNowBtn); expect(h.state.stop).toHaveBeenCalledTimes(1); expect(h.state.sendMessage).not.toHaveBeenCalled(); // The abort we triggered reaches onFinish: the promoted head is flushed. act(() => { h.state.onFinish?.({ message: { id: "a", role: "assistant", parts: [] }, isAbort: true, isDisconnect: false, isError: false, }); }); expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" }); }); it("tags exactly the next send as interrupted (one-shot flag)", () => { renderThread(); fireEvent.click(screen.getByTestId("queue-btn")); fireEvent.click(screen.getByLabelText("Send now")); const prep = h.state.transport!.prepareSendMessagesRequest; // The send right after "send now" carries interrupted: true... expect(prep({ messages: [], body: {} }).body.interrupted).toBe(true); // ...and only that one (the flag is read-and-cleared). expect(prep({ messages: [], body: {} }).body.interrupted).toBe(false); }); it("sends immediately without an interrupt when not streaming", () => { h.state.status = "ready"; renderThread(); fireEvent.click(screen.getByTestId("queue-btn")); fireEvent.click(screen.getByLabelText("Send now")); // No turn to interrupt: sent straight away, no abort, not flagged. expect(h.state.stop).not.toHaveBeenCalled(); expect(h.state.sendMessage).toHaveBeenCalledWith({ text: "queued text" }); const prep = h.state.transport!.prepareSendMessagesRequest; expect(prep({ messages: [], body: {} }).body.interrupted).toBe(false); }); });