import { act, renderHook } from "@testing-library/react"; import type { ReactNode } from "react"; import { describe, expect, it, vi } from "vitest"; import { useNanobotStream } from "@/hooks/useNanobotStream"; import type { InboundEvent } from "@/lib/types"; import { ClientProvider } from "@/providers/ClientProvider"; const EMPTY_MESSAGES: import("@/lib/types").UIMessage[] = []; function fakeClient() { const handlers = new Map void>>(); return { client: { status: "open" as const, defaultChatId: null as string | null, onStatus: () => () => {}, onError: () => () => {}, onChat(chatId: string, h: (ev: InboundEvent) => void) { let set = handlers.get(chatId); if (!set) { set = new Set(); handlers.set(chatId, set); } set.add(h); return () => set!.delete(h); }, sendMessage: vi.fn(), newChat: vi.fn(), attach: vi.fn(), connect: vi.fn(), close: vi.fn(), updateUrl: vi.fn(), }, emit(chatId: string, ev: InboundEvent) { const set = handlers.get(chatId); set?.forEach((h) => h(ev)); }, }; } function wrap(client: ReturnType["client"]) { return function Wrapper({ children }: { children: ReactNode }) { return ( {children} ); }; } describe("useNanobotStream", () => { it("starts in streaming mode when history shows pending tool calls", () => { const fake = fakeClient(); const initialMessages = [{ id: "m1", role: "assistant" as const, content: "Using tools", createdAt: Date.now(), }]; const { result } = renderHook( () => useNanobotStream("chat-p", initialMessages, true), { wrapper: wrap(fake.client), }, ); expect(result.current.isStreaming).toBe(true); }); it("collapses consecutive tool_hint frames into one trace row", () => { const fake = fakeClient(); const { result } = renderHook(() => useNanobotStream("chat-t", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); act(() => { fake.emit("chat-t", { event: "message", chat_id: "chat-t", text: 'weather("get")', kind: "tool_hint", }); fake.emit("chat-t", { event: "message", chat_id: "chat-t", text: 'search "hk weather"', kind: "tool_hint", }); }); expect(result.current.messages).toHaveLength(1); expect(result.current.messages[0].kind).toBe("trace"); expect(result.current.messages[0].role).toBe("tool"); expect(result.current.messages[0].traces).toEqual([ 'weather("get")', 'search "hk weather"', ]); act(() => { fake.emit("chat-t", { event: "message", chat_id: "chat-t", text: "## Summary", }); }); expect(result.current.messages).toHaveLength(2); expect(result.current.messages[1].role).toBe("assistant"); expect(result.current.messages[1].kind).toBeUndefined(); }); it("attaches assistant media_urls to complete messages", () => { const fake = fakeClient(); const { result } = renderHook(() => useNanobotStream("chat-m", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); act(() => { fake.emit("chat-m", { event: "message", chat_id: "chat-m", text: "video ready", media_urls: [{ url: "/api/media/sig/payload", name: "demo.mp4" }], }); }); expect(result.current.messages).toHaveLength(1); expect(result.current.messages[0].media).toEqual([ { kind: "video", url: "/api/media/sig/payload", name: "demo.mp4" }, ]); }); it("keeps assistant buttons on complete messages", () => { const fake = fakeClient(); const { result } = renderHook(() => useNanobotStream("chat-q", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); act(() => { fake.emit("chat-q", { event: "message", chat_id: "chat-q", text: "How should I continue?\n\n1. Short answer\n2. Detailed answer", button_prompt: "How should I continue?", buttons: [["Short answer", "Detailed answer"]], }); }); expect(result.current.messages).toHaveLength(1); expect(result.current.messages[0].content).toBe("How should I continue?"); expect(result.current.messages[0].buttons).toEqual([ ["Short answer", "Detailed answer"], ]); }); it("keeps streaming alive across stream_end and completes on turn_end", () => { const fake = fakeClient(); const { result } = renderHook(() => useNanobotStream("chat-s", EMPTY_MESSAGES), { wrapper: wrap(fake.client), }); act(() => { fake.emit("chat-s", { event: "delta", chat_id: "chat-s", text: "Hello", }); }); expect(result.current.isStreaming).toBe(true); expect(result.current.messages[0]).toMatchObject({ role: "assistant", content: "Hello", isStreaming: true, }); act(() => { fake.emit("chat-s", { event: "stream_end", chat_id: "chat-s", }); }); expect(result.current.isStreaming).toBe(true); expect(result.current.messages[0].isStreaming).toBe(true); act(() => { fake.emit("chat-s", { event: "message", chat_id: "chat-s", text: "Hello world", }); }); expect(result.current.isStreaming).toBe(true); expect(result.current.messages.at(-1)).toMatchObject({ role: "assistant", content: "Hello world", }); act(() => { fake.emit("chat-s", { event: "turn_end", chat_id: "chat-s", }); }); expect(result.current.isStreaming).toBe(false); expect(result.current.messages.every((message) => !message.isStreaming)).toBe(true); }); });