// @vitest-environment jsdom import type { ComponentProps } from "react"; import { act, createElement, useEffect } from "react"; import { createRoot } from "react-dom/client"; import type { AssistantClient, AssistantRuntime, AttachmentAdapter, ChatModelAdapter, ThreadMessageLike, } from "@assistant-ui/react"; import { defineToolkit, INTERNAL, MessagePrimitive, ThreadPrimitive, useAui, useAuiEvent, useLocalRuntime, } from "@assistant-ui/react"; import { LexicalComposerInput } from "@assistant-ui/react-lexical"; import { afterEach, describe, expect, test, vi } from "vitest"; import { z } from "zod/v4"; import { showErrorNotification } from "@homarr/notifications"; import { AssistantComposerSurfaceProvider, AssistantRunFocusPreserver, AssistantRuntimeProviderWithTools, assistantSurfaceComposerCacheLimit, } from "./assistant-runtime-provider"; vi.mock("@homarr/notifications", () => ({ showErrorNotification: vi.fn() })); vi.mock("@homarr/translation/client", () => ({ useScopedI18n: () => (key: string) => key })); (globalThis as typeof globalThis & { IS_REACT_ACT_ENVIRONMENT: boolean }).IS_REACT_ACT_ENVIRONMENT = true; const noOpAdapter: ChatModelAdapter = { async *run() {} }; const createAttachmentAdapter = (accept: string): AttachmentAdapter => ({ accept, async add({ file }) { return { id: `${file.name}-${file.size}`, type: file.type.startsWith("image/") ? "image" : "document", name: file.name, contentType: file.type, file, status: { type: "requires-action", reason: "composer-send" }, }; }, async send(attachment) { return { ...attachment, status: { type: "complete" }, content: [{ type: "text", text: attachment.name }], }; }, async remove() {}, }); const toolkit = defineToolkit({ ask_user: { type: "human", display: "standalone", parameters: z.object({ question: z.string() }), render: ({ args }) => createElement("div", { "data-testid": "ask-user-ui" }, `Question: ${args.question}`), }, }); const messages: ThreadMessageLike[] = [ { role: "assistant", content: [ { type: "tool-call", toolCallId: "ask-1", toolName: "ask_user", args: { question: "Which board?" }, argsText: JSON.stringify({ question: "Which board?" }), }, ], status: { type: "requires-action", reason: "tool-calls" }, }, ]; const Message = () => createElement(MessagePrimitive.Parts, { components: { tools: { Fallback: () => createElement("div", { "data-testid": "fallback-tool-ui" }, "Fallback"), }, }, }); const TestAssistant = () => { const runtime = useLocalRuntime(noOpAdapter, { initialMessages: messages }); return createElement( AssistantRuntimeProviderWithTools, { runtime, toolkit }, createElement(ThreadPrimitive.Messages, { components: { Message } }), ); }; interface SurfaceProbeProps { id: string; onReady: (id: string, aui: AssistantClient) => void; } const SurfaceProbe = ({ id, onReady }: SurfaceProbeProps) => { const aui = useAui(); useEffect(() => { onReady(id, aui); }, [aui, id, onReady]); return createElement(LexicalComposerInput, { "data-assistant-composer-input": true, "data-testid": `input-${id}`, } as ComponentProps); }; const LexicalRunFocusProbe = () => { useAuiEvent("thread.runStart", () => { queueMicrotask(() => { document.querySelector('[data-testid="input-b"] [contenteditable="true"]')?.focus(); }); }); return null; }; const SurfaceTestAssistant = ({ onReady, onRuntimeReady, allowImages = true, simulateRunFocus = false, }: { onReady: SurfaceProbeProps["onReady"]; onRuntimeReady: (runtime: AssistantRuntime) => void; allowImages?: boolean; simulateRunFocus?: boolean; }) => { const runtime = useLocalRuntime(noOpAdapter, { adapters: { attachments: createAttachmentAdapter(allowImages ? "image/png,text/plain" : "text/plain") }, }); useEffect(() => onRuntimeReady(runtime), [onRuntimeReady, runtime]); return createElement( AssistantRuntimeProviderWithTools, { runtime, toolkit }, createElement(AssistantRunFocusPreserver), createElement( AssistantComposerSurfaceProvider, { surfaceId: "surface-a" }, createElement(SurfaceProbe, { id: "a", onReady }), ), createElement( AssistantComposerSurfaceProvider, { surfaceId: "surface-b" }, createElement(SurfaceProbe, { id: "b", onReady }), ), simulateRunFocus ? createElement(LexicalRunFocusProbe) : null, ); }; const renderSurfaceTestAssistant = async ({ allowImages = true, simulateRunFocus = false, }: { allowImages?: boolean; simulateRunFocus?: boolean } = {}) => { const container = document.createElement("div"); containers.push(container); document.body.append(container); const root = createRoot(container); const clients = new Map(); let sharedRuntime: AssistantRuntime | undefined; await act(async () => { root.render( createElement(SurfaceTestAssistant, { allowImages, simulateRunFocus, onReady: (id, aui) => clients.set(id, aui), onRuntimeReady: (runtime) => { sharedRuntime = runtime; }, }), ); await new Promise((resolve) => setTimeout(resolve, 0)); }); const a = clients.get("a"); const b = clients.get("b"); if (!a || !b || !sharedRuntime) throw new Error("The assistant runtime and both surfaces must be registered"); return { container, root, a, b, sharedRuntime, }; }; const getRequiredComposerInput = (container: HTMLElement, testId: string) => { const wrapper = container.querySelector(`[data-testid="${testId}"]`); const target = wrapper?.querySelector('[contenteditable="true"]'); if (!target) throw new Error(`Missing composer input ${testId}`); return target; }; const getUserMessageText = (aui: AssistantClient) => aui .thread() .getState() .messages.filter((message) => message.role === "user") .flatMap((message) => message.content) .filter((part) => part.type === "text") .map((part) => part.text); const containers: HTMLDivElement[] = []; afterEach(() => { for (const container of containers.splice(0)) container.remove(); }); describe("AssistantRuntimeProviderWithTools", () => { test("makes the registered human-tool renderer available to message parts", async () => { const container = document.createElement("div"); containers.push(container); document.body.append(container); const root = createRoot(container); await act(async () => { root.render(createElement(TestAssistant)); await new Promise((resolve) => setTimeout(resolve, 0)); }); expect(container.querySelector('[data-testid="ask-user-ui"]')?.textContent).toBe("Question: Which board?"); expect(container.querySelector('[data-testid="fallback-tool-ui"]')).toBeNull(); await act(async () => root.unmount()); }); test("keeps drafts and attachments independent while both surfaces share one thread", async () => { const { container, root, a, b } = await renderSurfaceTestAssistant(); act(() => { a.composer().setText("Draft in the board widget"); b.composer().setText("Draft in the floating panel"); }); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); await act(async () => { await a.composer().addAttachment({ id: "board-file", name: "board.txt", contentType: "text/plain", content: [{ type: "text", text: "board attachment" }], }); }); expect(a.composer().getState().text).toBe("Draft in the board widget"); expect(b.composer().getState().text).toBe("Draft in the floating panel"); expect( a .composer() .getState() .attachments.map((attachment) => attachment.name), ).toEqual(["board.txt"]); expect(b.composer().getState().attachments).toHaveLength(0); expect(a.thread().getState().messages).toEqual(b.thread().getState().messages); expect(getRequiredComposerInput(container, "input-a").textContent).toBe("Draft in the board widget"); expect(getRequiredComposerInput(container, "input-b").textContent).toBe("Draft in the floating panel"); await act(async () => root.unmount()); }); test("sends either local composer into the same shared conversation without clearing the other draft", async () => { const { root, a, b } = await renderSurfaceTestAssistant(); act(() => { a.composer().setText("Message from the widget"); b.composer().setText("Message from the panel"); a.composer().send(); }); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); expect(getUserMessageText(a)).toContain("Message from the widget"); expect(getUserMessageText(b)).toContain("Message from the widget"); expect(a.composer().getState().text).toBe(""); expect(b.composer().getState().text).toBe("Message from the panel"); act(() => b.composer().send()); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); expect(getUserMessageText(a)).toContain("Message from the panel"); expect(getUserMessageText(b)).toContain("Message from the panel"); await act(async () => root.unmount()); }); test("keeps each surface draft scoped to the selected conversation", async () => { const { root, a, b } = await renderSurfaceTestAssistant(); act(() => { a.composer().setText("Initialize the first conversation"); a.composer().send(); }); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); const firstThreadId = a.threads().getState().mainThreadId; act(() => { a.composer().setText("First thread / widget"); b.composer().setText("First thread / panel"); }); await act(async () => a.threads().switchToNewThread()); expect(a.threads().getState().mainThreadId).not.toBe(firstThreadId); expect(a.composer().getState().text).toBe(""); expect(b.composer().getState().text).toBe(""); act(() => a.composer().setText("Second thread / widget")); await act(async () => a.threads().switchToThread(firstThreadId)); expect(a.composer().getState().text).toBe("First thread / widget"); expect(b.composer().getState().text).toBe("First thread / panel"); await act(async () => root.unmount()); }); test("forwards main-thread reloads through the surface runtime", async () => { const { root, a, sharedRuntime } = await renderSurfaceTestAssistant(); const reloadMainThread = vi.spyOn(sharedRuntime.threads, "reloadMainThread").mockResolvedValue(); await act(async () => a.threads().reloadMainThread()); expect(reloadMainThread).toHaveBeenCalledOnce(); await act(async () => root.unmount()); }); test("restores focus after the Lexical run-start behavior tries to focus every surface", async () => { const outsideInput = document.createElement("input"); document.body.append(outsideInput); const { root, a } = await renderSurfaceTestAssistant(); outsideInput.focus(); act(() => { a.composer().setText("Run without stealing focus"); a.composer().send(); }); await act(async () => { await new Promise((resolve) => setTimeout(resolve, 0)); }); expect(document.activeElement).toBe(outsideInput); outsideInput.remove(); await act(async () => root.unmount()); }); test("keeps focus and dashboard scroll on the surface that started the run", async () => { const { container, root, a } = await renderSurfaceTestAssistant({ simulateRunFocus: true }); const panelInput = getRequiredComposerInput(container, "input-a"); const widgetInput = getRequiredComposerInput(container, "input-b"); const scrollContainer = container; scrollContainer.scrollTop = 120; const focusWidget = widgetInput.focus.bind(widgetInput); const widgetFocus = vi.spyOn(widgetInput, "focus").mockImplementation((options) => { // jsdom does not implement native focus scrolling, so reproduce the dashboard jump that a // browser performs when the off-screen widget composer receives programmatic focus. scrollContainer.scrollTop = 640; focusWidget(options); }); panelInput.focus(); act(() => { a.composer().setText("Send from the floating panel"); a.composer().send(); }); await act(async () => { await new Promise((resolve) => setTimeout(resolve, 0)); }); expect(widgetFocus).toHaveBeenCalled(); expect(document.activeElement).toBe(panelInput); expect(scrollContainer.scrollTop).toBe(120); await act(async () => root.unmount()); }); test("adds pasted images only to the active surface composer", async () => { const { container, root, a, b } = await renderSurfaceTestAssistant(); const target = getRequiredComposerInput(container, "input-a"); const paste = new Event("paste", { bubbles: true, cancelable: true }); Object.defineProperty(paste, "clipboardData", { value: { files: [new File(["image"], "dashboard.png", { type: "image/png" })] }, }); await act(async () => { target.dispatchEvent(paste); await new Promise((resolve) => setTimeout(resolve, 0)); }); expect(paste.defaultPrevented).toBe(true); expect( a .composer() .getState() .attachments.map((attachment) => attachment.name), ).toEqual(["dashboard.png"]); expect(b.composer().getState().attachments).toHaveLength(0); await act(async () => root.unmount()); }); test("rejects a pasted image with an actionable message when the selected model is text-only", async () => { vi.mocked(showErrorNotification).mockClear(); const { container, root, a } = await renderSurfaceTestAssistant({ allowImages: false }); const target = getRequiredComposerInput(container, "input-a"); const paste = new Event("paste", { bubbles: true, cancelable: true }); Object.defineProperty(paste, "clipboardData", { value: { files: [new File(["image"], "dashboard.png", { type: "image/png" })] }, }); await act(async () => target.dispatchEvent(paste)); expect(a.composer().getState().attachments).toHaveLength(0); expect(showErrorNotification).toHaveBeenCalledWith({ title: "attachments.errorTitle", message: "attachments.imageUnsupported", }); await act(async () => root.unmount()); }); test("does not capture file paste from a non-composer textarea", async () => { const { container, root, a } = await renderSurfaceTestAssistant(); const textarea = document.createElement("textarea"); container.querySelector('[data-testid="input-a"]')?.parentElement?.append(textarea); const paste = new Event("paste", { bubbles: true, cancelable: true }); Object.defineProperty(paste, "clipboardData", { value: { files: [new File(["image"], "dashboard.png", { type: "image/png" })] }, }); await act(async () => textarea.dispatchEvent(paste)); expect(paste.defaultPrevented).toBe(false); expect(a.composer().getState().attachments).toHaveLength(0); await act(async () => root.unmount()); }); test("enforces the attachment limit independently for each surface", async () => { vi.mocked(showErrorNotification).mockClear(); const { root, a, b } = await renderSurfaceTestAssistant(); await act(async () => { await Promise.all( Array.from({ length: 5 }, (_, index) => a.composer().addAttachment(new File([`a-${index}`], `a-${index}.txt`, { type: "text/plain" })), ), ); await Promise.all( Array.from({ length: 5 }, (_, index) => b.composer().addAttachment(new File([`b-${index}`], `b-${index}.txt`, { type: "text/plain" })), ), ); }); expect(a.composer().getState().attachments).toHaveLength(5); expect(b.composer().getState().attachments).toHaveLength(5); await expect(a.composer().addAttachment(new File(["six"], "a-six.txt", { type: "text/plain" }))).rejects.toThrow( "A message can include up to 5 attachments.", ); expect(b.composer().getState().attachments).toHaveLength(5); await act(async () => root.unmount()); }); test("disposes scoped composer subscriptions when its surfaces unmount", async () => { const originalConnect = INTERNAL.DefaultThreadComposerRuntimeCore.prototype.connect; let activeConnections = 0; const connectSpy = vi .spyOn(INTERNAL.DefaultThreadComposerRuntimeCore.prototype, "connect") .mockImplementation(function (this: InstanceType) { const unsubscribe = originalConnect.call(this); activeConnections += 1; let connected = true; return () => { if (!connected) return; connected = false; activeConnections -= 1; unsubscribe(); }; }); const { root } = await renderSurfaceTestAssistant(); const mountedConnections = activeConnections; expect(mountedConnections).toBeGreaterThanOrEqual(3); await act(async () => root.unmount()); expect(activeConnections).toBeLessThan(mountedConnections); expect(mountedConnections - activeConnections).toBe(2); connectSpy.mockRestore(); }); test("recreates scoped composers when the shared conversation core is replaced", async () => { const originalConnect = INTERNAL.DefaultThreadComposerRuntimeCore.prototype.connect; let connectCalls = 0; let disconnectCalls = 0; const connectSpy = vi .spyOn(INTERNAL.DefaultThreadComposerRuntimeCore.prototype, "connect") .mockImplementation(function (this: InstanceType) { connectCalls += 1; const unsubscribe = originalConnect.call(this); let connected = true; return () => { if (!connected) return; connected = false; disconnectCalls += 1; unsubscribe(); }; }); const { root, a, b, sharedRuntime } = await renderSurfaceTestAssistant(); const mountedConnectCalls = connectCalls; act(() => { a.composer().setText("Draft on the first board core"); b.composer().setText("Draft on the first panel core"); }); const sharedThread = sharedRuntime.thread as InstanceType; // oxlint-disable-next-line no-underscore-dangle -- replacing the bound core reproduces assistant-ui's outer runtime swap. const binding = sharedThread.__internal_threadBinding; const originalGetState = binding.getState; const originalCore = originalGetState(); const replacementCore = new Proxy(originalCore, { get(target, property) { const value = Reflect.get(target, property, target) as unknown; return typeof value === "function" ? value.bind(target) : value; }, }); binding.getState = () => replacementCore; act(() => originalCore.reset()); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); expect(connectCalls - mountedConnectCalls).toBeGreaterThanOrEqual(2); expect(disconnectCalls).toBeGreaterThanOrEqual(2); expect(a.composer().getState().text).toBe("Draft on the first board core"); expect(b.composer().getState().text).toBe("Draft on the first panel core"); await act(async () => root.unmount()); binding.getState = originalGetState; connectSpy.mockRestore(); }); test("evicts the least recently used inactive conversation draft", async () => { const { root, a } = await renderSurfaceTestAssistant(); act(() => { a.composer().setText("Initialize the oldest conversation"); a.composer().send(); }); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); const oldestThreadId = a.threads().getState().mainThreadId; act(() => a.composer().setText("Draft that should eventually be evicted")); for (let index = 0; index < assistantSurfaceComposerCacheLimit; index += 1) { await act(async () => a.threads().switchToNewThread()); act(() => { a.composer().setText(`Initialize conversation ${index}`); a.composer().send(); }); await act(async () => new Promise((resolve) => setTimeout(resolve, 0))); } await act(async () => a.threads().switchToThread(oldestThreadId)); expect(a.composer().getState().text).toBe(""); await act(async () => root.unmount()); }); });