@assistant-ui/core 0.3.12 → 0.3.13

package/src/tests/external-store-thread-runtime-core-adapter.test.ts

@@ -6,6 +6,8 @@ import {
 import type { ExternalStoreAdapter } from "../runtimes/external-store/external-store-adapter";
 import type { ModelContextProvider } from "../model-context/types";
 import type { AppendMessage, ThreadMessage } from "../types/message";
+import { createMessageQueue } from "../runtime/queue/message-queue";
+import { getThreadMessageText } from "../utils/text";
 
 const createContextProvider = (): ModelContextProvider => ({
   getModelContext: () => ({}),
@@ -235,6 +237,175 @@ describe("ExternalStoreThreadRuntimeCore adapter contract", () => {
       expect(lastCall.map((m) => m.id)).toContain("server-msg");
     });
 
+    it("does not revert a store update that lands before the resync flushes", async () => {
+      const optimisticAssistant = {
+        ...createAssistantMessage("server-msg", "partial answer"),
+        status: { type: "running" as const },
+        metadata: {
+          unstable_state: null,
+          unstable_annotations: [],
+          unstable_data: [],
+          steps: [],
+          custom: {},
+          isOptimistic: true,
+        },
+      } as ThreadMessage;
+      const setMessages = vi.fn();
+      const adapter = createBaseAdapter({
+        messages: [createUserMessage("u1"), optimisticAssistant],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages,
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+
+      // The store settles the cancelled turn and re-supplies it in the same
+      // tick as the cancel.
+      core.__internal_setAdapter(
+        createBaseAdapter({
+          messages: [
+            createUserMessage("u1"),
+            createAssistantMessage("server-msg", "partial answer (stopped)"),
+          ],
+          isRunning: false,
+          onCancel: vi.fn(),
+          setMessages,
+        }),
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(setMessages).toHaveBeenCalled();
+      const lastCall = setMessages.mock.lastCall?.[0] as ThreadMessage[];
+      const texts = lastCall.map(getThreadMessageText);
+      expect(texts).toContain("partial answer (stopped)");
+      expect(texts).not.toContain("partial answer");
+    });
+
+    it("re-applies the user leaf rollback when the store updates before the flush", async () => {
+      const messages = [createUserMessage("u1", "cancel me")];
+      const setMessages = vi.fn();
+      const adapter = createBaseAdapter({
+        messages,
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages,
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+
+      // The host flips isRunning after onCancel; the store still holds the
+      // rolled-back message because only the deferred resync removes it.
+      core.__internal_setAdapter(
+        createBaseAdapter({
+          messages,
+          isRunning: false,
+          onCancel: vi.fn(),
+          setMessages,
+        }),
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("cancel me");
+      expect(setMessages).toHaveBeenCalled();
+      const lastCall = setMessages.mock.lastCall?.[0] as ThreadMessage[];
+      expect(lastCall.map((m) => m.id)).not.toContain("u1");
+      expect(core.export().messages.map((m) => m.message.id)).not.toContain(
+        "u1",
+      );
+    });
+
+    it("keeps a rolled-back user leaf the store answered in the gap and takes the draft back", async () => {
+      const setMessages = vi.fn();
+      const adapter = createBaseAdapter({
+        messages: [createUserMessage("u1", "cancel me")],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages,
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+      expect(core.composer.text).toBe("cancel me");
+
+      core.__internal_setAdapter(
+        createBaseAdapter({
+          messages: [
+            createUserMessage("u1", "cancel me"),
+            createAssistantMessage("a1", "answered anyway"),
+          ],
+          isRunning: false,
+          onCancel: vi.fn(),
+          setMessages,
+        }),
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("");
+      const lastCall = setMessages.mock.lastCall?.[0] as ThreadMessage[];
+      expect(lastCall.map((m) => m.id)).toEqual(["u1", "a1"]);
+    });
+
+    it("leaves an edited draft alone when the store answered in the gap", async () => {
+      const setMessages = vi.fn();
+      const adapter = createBaseAdapter({
+        messages: [createUserMessage("u1", "cancel me")],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages,
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+      core.composer.setText("edited");
+
+      core.__internal_setAdapter(
+        createBaseAdapter({
+          messages: [
+            createUserMessage("u1", "cancel me"),
+            createAssistantMessage("a1", "answered anyway"),
+          ],
+          isRunning: false,
+          onCancel: vi.fn(),
+          setMessages,
+        }),
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("edited");
+      const lastCall = setMessages.mock.lastCall?.[0] as ThreadMessage[];
+      expect(lastCall.map((m) => m.id)).toEqual(["u1", "a1"]);
+    });
+
+    it("drops a placeholder regenerated between cancel and flush", async () => {
+      const setMessages = vi.fn();
+      const adapter = createBaseAdapter({
+        messages: [createUserMessage("u1", "cancel me")],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages,
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+
+      core.__internal_setAdapter(
+        createBaseAdapter({
+          messages: [createUserMessage("u1", "cancel me")],
+          isRunning: true,
+          onCancel: vi.fn(),
+          setMessages,
+        }),
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("cancel me");
+      const lastCall = setMessages.mock.lastCall?.[0] as ThreadMessage[];
+      expect(lastCall).toEqual([]);
+    });
+
     it("evicts an empty optimistic head on cancel", async () => {
       const optimisticAssistant = {
         ...createAssistantMessage("server-msg", ""),
@@ -264,6 +435,135 @@ describe("ExternalStoreThreadRuntimeCore adapter contract", () => {
       const lastCall = setMessages.mock.lastCall?.[0] as ThreadMessage[];
       expect(lastCall.map((m) => m.id)).not.toContain("server-msg");
     });
+
+    it("pauses a queue instead of dispatching the next message", async () => {
+      const dispatched: string[] = [];
+      const queue = createMessageQueue({
+        run: (message) => {
+          dispatched.push(getThreadMessageText(message));
+        },
+      });
+      const adapter = createBaseAdapter({
+        isRunning: true,
+        onCancel: vi.fn(),
+        queue: queue.adapter,
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+      const appendAtTail = (text: string) =>
+        core.append({
+          role: "user",
+          content: [{ type: "text", text }],
+          attachments: [],
+          createdAt: new Date(),
+          parentId: core.messages.at(-1)?.id ?? null,
+          sourceId: null,
+          runConfig: undefined,
+          metadata: { custom: {} },
+        } as AppendMessage);
+
+      await appendAtTail("first");
+      await appendAtTail("queued");
+
+      core.cancelRun();
+      // the aborted run settles; a queue host reports that the same way it
+      // reports any other run ending
+      queue.notifyIdle();
+
+      expect(dispatched).toEqual(["first"]);
+      expect(queue.adapter.steerItems.map((item) => item.prompt)).toEqual([
+        "queued",
+      ]);
+    });
+
+    it("keeps a trailing user message out of the composer without setMessages", async () => {
+      const userMessage = createUserMessage("u1", "Do not duplicate me");
+      const adapter = createBaseAdapter({
+        messages: [userMessage],
+        isRunning: true,
+        onCancel: vi.fn(),
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("");
+      expect(core.export().messages.map((m) => m.message.id)).toContain("u1");
+    });
+
+    it("keeps the message in the thread when the composer is busy", async () => {
+      const adapter = createBaseAdapter({
+        messages: [createUserMessage("u1", "cancel me")],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages: vi.fn(),
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+      core.composer.setText("something else");
+
+      core.cancelRun();
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("something else");
+      expect(core.export().messages.map((m) => m.message.id)).toContain("u1");
+    });
+
+    it("keeps a message with non-text content in the thread", async () => {
+      const userMessage = {
+        ...createUserMessage("u1", "look at this"),
+        content: [
+          { type: "text" as const, text: "look at this" },
+          { type: "image" as const, image: "https://example.com/cat.png" },
+        ],
+      } as ThreadMessage;
+      const adapter = createBaseAdapter({
+        messages: [userMessage],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages: vi.fn(),
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("");
+      expect(core.export().messages.map((m) => m.message.id)).toContain("u1");
+    });
+
+    it("restores the message whole when the composer is free", async () => {
+      const attachment = {
+        id: "a1",
+        type: "file" as const,
+        name: "spec.pdf",
+        contentType: "application/pdf",
+        status: { type: "complete" as const },
+        content: [],
+      };
+      const quote = { text: "quoted", messageId: "m-1" };
+      const userMessage = {
+        ...createUserMessage("u1", "cancel me"),
+        attachments: [attachment],
+        metadata: { custom: { quote } },
+      } as ThreadMessage;
+      const adapter = createBaseAdapter({
+        messages: [userMessage],
+        isRunning: true,
+        onCancel: vi.fn(),
+        setMessages: vi.fn(),
+      });
+      const core = new ExternalStoreThreadRuntimeCore(contextProvider, adapter);
+
+      core.cancelRun();
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(core.composer.text).toBe("cancel me");
+      expect(core.composer.attachments).toEqual([attachment]);
+      expect(core.composer.quote).toEqual(quote);
+      expect(core.export().messages.map((m) => m.message.id)).not.toContain(
+        "u1",
+      );
+    });
   });
 
   describe("optimistic assistant message", () => {

package/src/tests/external-thread-attachments.test.tsx

@@ -9,6 +9,10 @@ import type {
   ExternalThreadProps,
 } from "../store/clients/external-thread";
 import { ExternalThread } from "../store/clients/external-thread";
+import type {
+  CompleteAttachment,
+  PendingAttachment,
+} from "../types/attachment";
 
 const renderThreadWithProps = (props: Partial<ExternalThreadProps>) => {
   const captured: { aui?: ReturnType<typeof useAui> } = {};
@@ -215,6 +219,59 @@ describe("ExternalThread attachments", () => {
     );
   });
 
+  it("preserves composer state while attachments are prepared for send", async () => {
+    let resolveSend!: (attachment: CompleteAttachment) => void;
+    const onNew = vi.fn();
+    const file = new File(["data"], "notes.txt", { type: "text/plain" });
+    const adapter = {
+      accept: "*",
+      add: async () => ({
+        id: "att-1",
+        type: "file" as const,
+        name: file.name,
+        contentType: file.type,
+        file,
+        status: {
+          type: "requires-action" as const,
+          reason: "composer-send" as const,
+        },
+      }),
+      send: () =>
+        new Promise<CompleteAttachment>((resolve) => {
+          resolveSend = resolve;
+        }),
+      remove: async () => {},
+    };
+    const aui = renderThreadWithProps({ attachmentAdapter: adapter, onNew });
+    const composer = () => aui().thread.composer();
+
+    await act(() => composer().addAttachment(file));
+    act(() => {
+      composer().setText("first message");
+      composer().setRole("assistant");
+      composer().setRunConfig({ model: "model-a" });
+      composer().send();
+      composer().setRole("system");
+      composer().setRunConfig({ model: "model-b" });
+    });
+    await act(async () => {
+      resolveSend({
+        id: "att-1",
+        type: "file",
+        name: file.name,
+        contentType: file.type,
+        status: { type: "complete" },
+        content: [],
+      });
+    });
+
+    await waitFor(() => expect(onNew).toHaveBeenCalledTimes(1));
+    expect(onNew.mock.calls[0]![0]).toMatchObject({
+      role: "assistant",
+      runConfig: { model: "model-a" },
+    });
+  });
+
   it.each([
     [
       "clearAttachments",
@@ -263,6 +320,163 @@ describe("ExternalThread attachments", () => {
     },
   );
 
+  it.each([
+    [
+      "clearAttachments",
+      (c: { clearAttachments(): Promise<void> }) => c.clearAttachments(),
+    ],
+    ["reset", (c: { reset(): Promise<void> }) => c.reset()],
+  ] as const)(
+    "does not restore an attachment that resolves after %s",
+    async (_name, invoke) => {
+      let resolveAdd!: (attachment: PendingAttachment) => void;
+      const add = vi.fn(
+        () =>
+          new Promise<PendingAttachment>((resolve) => {
+            resolveAdd = resolve;
+          }),
+      );
+      const aui = renderThreadWithProps({
+        attachmentAdapter: {
+          accept: "*",
+          add,
+          send: async () => ({}) as never,
+          remove: async () => {},
+        },
+      });
+      const file = new File(["data"], "notes.txt", { type: "text/plain" });
+
+      let addPromise!: Promise<void>;
+      act(() => {
+        addPromise = aui().thread.composer().addAttachment(file);
+      });
+      await act(() => invoke(aui().thread.composer()));
+      await act(async () => {
+        resolveAdd({
+          id: "att-1",
+          type: "file",
+          name: file.name,
+          contentType: file.type,
+          file,
+          status: { type: "requires-action", reason: "composer-send" },
+        });
+        await addPromise;
+      });
+
+      expect(aui().thread.composer().getState().attachments).toHaveLength(0);
+    },
+  );
+
+  it("stops attachment generator updates after removal", async () => {
+    let resumeAdd!: () => void;
+    const drainedAfterRemoval = vi.fn();
+    const file = new File(["data"], "notes.txt", { type: "text/plain" });
+    const pendingAttachment: PendingAttachment = {
+      id: "att-1",
+      type: "file",
+      name: file.name,
+      contentType: file.type,
+      file,
+      status: { type: "requires-action", reason: "composer-send" },
+    };
+    const add = vi.fn(async function* () {
+      yield pendingAttachment;
+      await new Promise<void>((resolve) => {
+        resumeAdd = resolve;
+      });
+      yield {
+        ...pendingAttachment,
+        status: { type: "running", reason: "uploading", progress: 1 },
+      } satisfies PendingAttachment;
+      drainedAfterRemoval();
+    });
+    const aui = renderThreadWithProps({
+      attachmentAdapter: {
+        accept: "*",
+        add,
+        send: async () => ({}) as never,
+        remove: async () => {},
+      },
+    });
+    const composer = () => aui().thread.composer();
+
+    let addPromise!: Promise<void>;
+    act(() => {
+      addPromise = composer().addAttachment(file);
+    });
+    await waitFor(() =>
+      expect(composer().getState().attachments).toHaveLength(1),
+    );
+    await act(() => composer().attachment({ id: "att-1" }).remove());
+    await act(async () => {
+      resumeAdd();
+      await addPromise;
+    });
+
+    expect(composer().getState().attachments).toHaveLength(0);
+    expect(drainedAfterRemoval).not.toHaveBeenCalled();
+  });
+
+  it("does not restore attachment generator updates after send", async () => {
+    let resumeAdd!: () => void;
+    const drainedAfterSend = vi.fn();
+    const onNew = vi.fn();
+    const file = new File(["data"], "notes.txt", { type: "text/plain" });
+    const pendingAttachment: PendingAttachment = {
+      id: "att-1",
+      type: "file",
+      name: file.name,
+      contentType: file.type,
+      file,
+      status: { type: "requires-action", reason: "composer-send" },
+    };
+    const add = vi.fn(async function* () {
+      yield pendingAttachment;
+      await new Promise<void>((resolve) => {
+        resumeAdd = resolve;
+      });
+      yield {
+        ...pendingAttachment,
+        status: { type: "running", reason: "uploading", progress: 1 },
+      } satisfies PendingAttachment;
+      drainedAfterSend();
+    });
+    const aui = renderThreadWithProps({
+      attachmentAdapter: {
+        accept: "*",
+        add,
+        send: async (attachment) => ({
+          ...attachment,
+          status: { type: "complete" },
+          content: [],
+        }),
+        remove: async () => {},
+      },
+      onNew,
+    });
+    const composer = () => aui().thread.composer();
+
+    let addPromise!: Promise<void>;
+    act(() => {
+      addPromise = composer().addAttachment(file);
+    });
+    await waitFor(() =>
+      expect(composer().getState().attachments).toHaveLength(1),
+    );
+    act(() => {
+      composer().setText("hello");
+      composer().send();
+    });
+    await waitFor(() => expect(onNew).toHaveBeenCalledTimes(1));
+    await act(async () => {
+      resumeAdd();
+      await addPromise;
+    });
+
+    expect(composer().getState().attachments).toHaveLength(0);
+    expect(drainedAfterSend).not.toHaveBeenCalled();
+  });
+
   it("routes edit-composer attachments through the adapter", async () => {
     const add = vi.fn(async ({ file }: { file: File }) => ({
       id: "att-edit",

package/src/tests/external-thread-parity.test.tsx

@@ -279,6 +279,60 @@ describe("ExternalThread composer", () => {
     expect(enqueue).not.toHaveBeenCalled();
   });
 
+  it("pauses the queue on cancel", async () => {
+    const notifyCancelled = vi.fn();
+    const onCancel = vi.fn();
+    const { aui } = renderThread({
+      messages: [],
+      isRunning: true,
+      onCancel,
+      queue: {
+        items: [],
+        steerItems: [],
+        enqueue: vi.fn(),
+        steer: vi.fn(),
+        move: vi.fn(),
+        edit: vi.fn(),
+        remove: vi.fn(),
+        __internal_notifyCancelled: notifyCancelled,
+      },
+    });
+
+    aui().thread.cancelRun();
+
+    expect(notifyCancelled).toHaveBeenCalledTimes(1);
+    expect(onCancel).toHaveBeenCalledTimes(1);
+    expect(notifyCancelled.mock.invocationCallOrder[0]!).toBeLessThan(
+      onCancel.mock.invocationCallOrder[0]!,
+    );
+    expect(aui().thread.getState().capabilities.cancel).toBe(true);
+    expect(aui().thread.composer().getState().canCancel).toBe(true);
+  });
+
+  it("leaves the queue alone when the host cannot cancel", async () => {
+    const notifyCancelled = vi.fn();
+    const { aui } = renderThread({
+      messages: [],
+      isRunning: true,
+      queue: {
+        items: [],
+        steerItems: [],
+        enqueue: vi.fn(),
+        steer: vi.fn(),
+        move: vi.fn(),
+        edit: vi.fn(),
+        remove: vi.fn(),
+        __internal_notifyCancelled: notifyCancelled,
+      },
+    });
+
+    aui().thread.cancelRun();
+
+    expect(notifyCancelled).not.toHaveBeenCalled();
+    expect(aui().thread.getState().capabilities.cancel).toBe(false);
+    expect(aui().thread.composer().getState().canCancel).toBe(false);
+  });
+
   it("routes edit-composer sends to onEdit with sourceId, bypassing the queue", async () => {
     const onEdit = vi.fn();
     const enqueue = vi.fn();