import { describe, it, expect, beforeEach, vi } from "vitest"; import { ExternalStoreThreadRuntimeCore, hasUpcomingMessage, } from "../legacy-runtime/runtime-cores/external-store/ExternalStoreThreadRuntimeCore"; import type { ExternalStoreAdapter } from "../legacy-runtime/runtime-cores/external-store/ExternalStoreAdapter"; import type { ModelContextProvider } from "@assistant-ui/core"; import type { AppendMessage, ThreadMessage } from "@assistant-ui/core"; const createContextProvider = (): ModelContextProvider => ({ getModelContext: () => ({}), }); const createUserMessage = (id: string, text = "Hello"): ThreadMessage => ({ id, role: "user" as const, createdAt: new Date(), content: [{ type: "text" as const, text }], attachments: [], metadata: { custom: {}, }, }) as ThreadMessage; const createAssistantMessage = (id: string, text = "Hi there"): ThreadMessage => ({ id, role: "assistant" as const, createdAt: new Date(), content: [{ type: "text" as const, text }], status: { type: "complete" as const, reason: "stop" as const }, metadata: { unstable_state: null, unstable_annotations: [], unstable_data: [], steps: [], custom: {}, }, }) as ThreadMessage; const createBaseAdapter = ( overrides: Partial> = {}, ): ExternalStoreAdapter => ({ messages: [], onNew: vi.fn(async () => {}), ...overrides, }); describe("ExternalStoreThreadRuntimeCore", () => { let contextProvider: ModelContextProvider; beforeEach(() => { contextProvider = createContextProvider(); }); describe("hasUpcomingMessage", () => { it("returns true when running and last message is not assistant", () => { const messages = [createUserMessage("u1")]; expect(hasUpcomingMessage(true, messages)).toBe(true); }); it("returns false when running and last message is assistant", () => { const messages = [createAssistantMessage("a1")]; expect(hasUpcomingMessage(true, messages)).toBe(false); }); it("returns false when not running", () => { const messages = [createUserMessage("u1")]; expect(hasUpcomingMessage(false, messages)).toBe(false); }); it("returns true for empty messages when running", () => { expect(hasUpcomingMessage(true, [])).toBe(true); }); }); describe("capabilities derived from adapter", () => { it("enables edit when onEdit is provided", () => { const adapter = createBaseAdapter({ onEdit: vi.fn(async () => {}), }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.edit).toBe(true); }); it("disables edit when onEdit is not provided", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.edit).toBe(false); }); it("enables reload when onReload is provided", () => { const adapter = createBaseAdapter({ onReload: vi.fn(async () => {}), }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.reload).toBe(true); }); it("disables reload when onReload is not provided", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.reload).toBe(false); }); it("enables cancel when onCancel is provided", () => { const adapter = createBaseAdapter({ onCancel: vi.fn(async () => {}), }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.cancel).toBe(true); }); it("disables cancel when onCancel is not provided", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.cancel).toBe(false); }); it("enables switchToBranch when setMessages is provided", () => { const adapter = createBaseAdapter({ setMessages: vi.fn(), }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.switchToBranch).toBe(true); }); it("defaults unstable_copy to true", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.unstable_copy).toBe(true); }); it("disables unstable_copy when explicitly set to false", () => { const adapter = createBaseAdapter({ unstable_capabilities: { copy: false }, }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.capabilities.unstable_copy).toBe(false); }); }); describe("append", () => { it("calls onNew when parentId matches last message id", async () => { const onNew = vi.fn(async () => {}); const messages = [createUserMessage("u1"), createAssistantMessage("a1")]; const adapter = createBaseAdapter({ messages, onNew }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const appendMessage: AppendMessage = { role: "user", content: [{ type: "text", text: "Follow up" }], attachments: [], createdAt: new Date(), parentId: "a1", sourceId: null, runConfig: undefined, metadata: { custom: {} }, } as AppendMessage; await core.append(appendMessage); expect(onNew).toHaveBeenCalledWith(appendMessage); }); it("calls onEdit when parentId does not match last message id", async () => { const onEdit = vi.fn(async () => {}); const onNew = vi.fn(async () => {}); const messages = [createUserMessage("u1"), createAssistantMessage("a1")]; const adapter = createBaseAdapter({ messages, onNew, onEdit }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const appendMessage: AppendMessage = { role: "user", content: [{ type: "text", text: "Edit" }], attachments: [], createdAt: new Date(), parentId: "u1", sourceId: null, runConfig: undefined, metadata: { custom: {} }, } as AppendMessage; await core.append(appendMessage); expect(onEdit).toHaveBeenCalledWith(appendMessage); expect(onNew).not.toHaveBeenCalled(); }); it("throws when adapter has no onEdit and parentId differs from head", async () => { const messages = [createUserMessage("u1"), createAssistantMessage("a1")]; const adapter = createBaseAdapter({ messages }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const appendMessage: AppendMessage = { role: "user", content: [{ type: "text", text: "Edit" }], attachments: [], createdAt: new Date(), parentId: "u1", sourceId: null, runConfig: undefined, metadata: { custom: {} }, } as AppendMessage; await expect(core.append(appendMessage)).rejects.toThrow( "Runtime does not support editing messages.", ); }); }); describe("startRun", () => { it("delegates to onReload", async () => { const onReload = vi.fn(async () => {}); const adapter = createBaseAdapter({ onReload }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); await core.startRun({ parentId: "msg-1", sourceId: null, runConfig: {} }); expect(onReload).toHaveBeenCalledWith("msg-1", { parentId: "msg-1", sourceId: null, runConfig: {}, }); }); it("throws when adapter has no onReload", async () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); await expect( core.startRun({ parentId: "msg-1", sourceId: null, runConfig: {} }), ).rejects.toThrow("Runtime does not support reloading messages."); }); }); describe("cancelRun", () => { it("delegates to onCancel", () => { const onCancel = vi.fn(); const messages = [createUserMessage("u1"), createAssistantMessage("a1")]; const adapter = createBaseAdapter({ messages, onCancel, setMessages: vi.fn(), }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); core.cancelRun(); expect(onCancel).toHaveBeenCalled(); }); it("throws when adapter has no onCancel", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(() => core.cancelRun()).toThrow( "Runtime does not support cancelling runs.", ); }); }); describe("optimistic assistant message", () => { it("adds optimistic assistant message when running with last user message", () => { const userMsg = createUserMessage("u1"); const adapter = createBaseAdapter({ messages: [userMsg], isRunning: true, }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const messages = core.messages; expect(messages.length).toBe(2); expect(messages[0]!.id).toBe("u1"); expect(messages[1]!.role).toBe("assistant"); expect(messages[1]!.content).toEqual([]); }); it("does not add optimistic message when last message is assistant", () => { const adapter = createBaseAdapter({ messages: [createUserMessage("u1"), createAssistantMessage("a1")], isRunning: true, }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const messages = core.messages; expect(messages.length).toBe(2); expect(messages[1]!.role).toBe("assistant"); expect(messages[1]!.id).toBe("a1"); }); it("does not add optimistic message when not running", () => { const adapter = createBaseAdapter({ messages: [createUserMessage("u1")], isRunning: false, }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.messages.length).toBe(1); }); }); describe("adapter update via __internal_setAdapter", () => { it("updates capabilities when adapter changes", () => { const adapter1 = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore( contextProvider, adapter1, ); expect(core.capabilities.edit).toBe(false); const adapter2 = createBaseAdapter({ onEdit: vi.fn(async () => {}), }); core.__internal_setAdapter(adapter2); expect(core.capabilities.edit).toBe(true); }); it("updates isDisabled from adapter", () => { const adapter = createBaseAdapter({ isDisabled: true }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.isDisabled).toBe(true); core.__internal_setAdapter(createBaseAdapter({ isDisabled: false })); expect(core.isDisabled).toBe(false); }); it("updates suggestions from adapter", () => { const suggestions = [{ prompt: "Tell me a joke" }]; const adapter = createBaseAdapter({ suggestions }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.suggestions).toEqual(suggestions); }); it("skips update when same adapter reference is passed", () => { const adapter = createBaseAdapter({ messages: [createUserMessage("u1")], }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const spy = vi.fn(); core.subscribe(spy); spy.mockClear(); // Same reference — should not notify core.__internal_setAdapter(adapter); expect(spy).not.toHaveBeenCalled(); }); }); describe("switchToBranch", () => { it("throws when setMessages is not provided", () => { const adapter = createBaseAdapter({ messages: [createUserMessage("u1")], }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(() => core.switchToBranch("u1")).toThrow( "Runtime does not support switching branches.", ); }); it("silently ignores branch switch while running", () => { const setMessages = vi.fn(); const adapter = createBaseAdapter({ messages: [createUserMessage("u1"), createAssistantMessage("a1")], isRunning: true, setMessages, }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); // Should not throw, should silently return core.switchToBranch("u1"); expect(setMessages).not.toHaveBeenCalled(); }); }); describe("exportExternalState and importExternalState", () => { it("delegates export to adapter", () => { const state = { key: "value" }; const adapter = createBaseAdapter({ onExportExternalState: vi.fn(() => state), }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(core.exportExternalState()).toEqual(state); }); it("throws on export when adapter has no onExportExternalState", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(() => core.exportExternalState()).toThrow( "Runtime does not support exporting external states.", ); }); it("delegates import to adapter", () => { const onLoadExternalState = vi.fn(); const adapter = createBaseAdapter({ onLoadExternalState }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const state = { key: "value" }; core.importExternalState(state); expect(onLoadExternalState).toHaveBeenCalledWith(state); }); it("throws on import when adapter has no onLoadExternalState", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(() => core.importExternalState({})).toThrow( "Runtime does not support importing external states.", ); }); }); describe("beginEdit", () => { it("throws when adapter has no onEdit", () => { const adapter = createBaseAdapter({ messages: [createUserMessage("u1")], }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(() => core.beginEdit("u1")).toThrow( "Runtime does not support editing.", ); }); }); describe("subscribe", () => { it("notifies subscribers when adapter changes", () => { const adapter1 = createBaseAdapter({ messages: [createUserMessage("u1")], }); const core = new ExternalStoreThreadRuntimeCore( contextProvider, adapter1, ); const callback = vi.fn(); core.subscribe(callback); callback.mockClear(); const adapter2 = createBaseAdapter({ messages: [createUserMessage("u1"), createAssistantMessage("a1")], }); core.__internal_setAdapter(adapter2); expect(callback).toHaveBeenCalled(); }); it("returns unsubscribe function", () => { const adapter = createBaseAdapter({ messages: [createUserMessage("u1")], }); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); const callback = vi.fn(); const unsub = core.subscribe(callback); callback.mockClear(); unsub(); core.__internal_setAdapter( createBaseAdapter({ messages: [createUserMessage("u1"), createAssistantMessage("a1")], }), ); expect(callback).not.toHaveBeenCalled(); }); }); describe("addToolResult", () => { it("throws when adapter has no onAddToolResult", () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); expect(() => core.addToolResult({ messageId: "m1", toolName: "tool", toolCallId: "tc1", result: {}, isError: false, }), ).toThrow("Runtime does not support tool results."); }); }); describe("resumeRun", () => { it("throws when adapter has no onResume", async () => { const adapter = createBaseAdapter(); const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter); await expect( core.resumeRun({ parentId: "msg-1" } as any), ).rejects.toThrow("Runtime does not support resuming runs."); }); }); describe("adapter requires messages or messageRepository", () => { it("throws when adapter has neither messages nor messageRepository", () => { const adapter = { onNew: vi.fn(async () => {}), } as unknown as ExternalStoreAdapter; expect( () => new ExternalStoreThreadRuntimeCore(contextProvider, adapter), ).toThrow( "ExternalStoreAdapter must provide either 'messages' or 'messageRepository'", ); }); }); });