@assistant-ui/core 0.3.12 → 0.3.13

package/src/tests/append-interactable-snapshots.test.ts

@@ -0,0 +1,350 @@
+import { describe, expect, it, vi } from "vitest";
+import { ExternalStoreRuntimeCore } from "../runtimes/external-store/external-store-runtime-core";
+import { LocalRuntimeCore } from "../runtimes/local/local-runtime-core";
+import { createMessageQueue } from "../runtime/queue/message-queue";
+import { getThreadMessageText } from "../utils/text";
+import type { ChatModelAdapter } from "../runtime/utils/chat-model-adapter";
+import type { AppendMessage, ThreadMessage } from "../types/message";
+
+const live = (state: unknown) => ({
+  interactables: [{ id: "n1", name: "note", state }],
+});
+
+const snapshotMessage = (id: string, state: unknown) =>
+  ({
+    id,
+    role: "user",
+    createdAt: new Date(),
+    content: [{ type: "text", text: "old" }],
+    attachments: [],
+    metadata: {
+      custom: { interactables: [{ id: "n1", name: "note", state }] },
+    },
+  }) as ThreadMessage;
+
+const userMessage = (text: string, parentId: string | null): AppendMessage => ({
+  parentId,
+  sourceId: null,
+  runConfig: {},
+  role: "user",
+  content: [{ type: "text", text }],
+  attachments: [],
+  metadata: { custom: {} },
+  createdAt: new Date(),
+});
+
+const externalThread = (options?: {
+  messages?: ThreadMessage[];
+  composerMetadata?: Record<string, unknown>;
+}) => {
+  const onNew = vi.fn(async () => {});
+  const onEdit = vi.fn(async () => {});
+  const core = new ExternalStoreRuntimeCore({
+    messages: options?.messages ?? [],
+    onNew,
+    onEdit,
+  });
+  if (options?.composerMetadata) {
+    core.registerModelContextProvider({
+      getModelContext: () => ({
+        unstable_composerMetadata: options.composerMetadata,
+      }),
+    });
+  }
+  return { thread: core.threads.getMainThreadRuntimeCore(), onNew, onEdit };
+};
+
+const appendedSnapshots = (sink: ReturnType<typeof vi.fn>) =>
+  (sink.mock.calls[0]![0] as AppendMessage).metadata?.custom?.interactables;
+
+const tailId = (messages: ThreadMessage[]) => messages.at(-1)?.id ?? null;
+
+describe("interactable snapshots on the append seam", () => {
+  it("stamps a baseline on the first send of the thread", async () => {
+    const { thread, onNew } = externalThread({
+      composerMetadata: live({ v: 1 }),
+    });
+    await thread.append(userMessage("hello", null));
+    expect(appendedSnapshots(onNew)).toEqual([
+      { id: "n1", name: "note", state: { v: 1 } },
+    ]);
+  });
+
+  it("stamps nothing when the history already carries the live state", async () => {
+    const messages = [snapshotMessage("msg-1", { v: 1 })];
+    const { thread, onNew } = externalThread({
+      messages,
+      composerMetadata: live({ v: 1 }),
+    });
+    await thread.append(userMessage("hello", tailId(messages)));
+    expect(appendedSnapshots(onNew)).toBeUndefined();
+  });
+
+  it("stamps the new state when it diverged from the latest snapshot", async () => {
+    const messages = [snapshotMessage("msg-1", { v: 1 })];
+    const { thread, onNew } = externalThread({
+      messages,
+      composerMetadata: live({ v: 2 }),
+    });
+    await thread.append(userMessage("hello", tailId(messages)));
+    expect(appendedSnapshots(onNew)).toEqual([
+      { id: "n1", name: "note", state: { v: 2 } },
+    ]);
+  });
+
+  it("stamps a programmatic append, the path suggestions and thread.append() take", async () => {
+    const { thread, onNew } = externalThread({
+      composerMetadata: live({ tasks: [] }),
+    });
+    await thread.append(userMessage("Add 3 tasks for a grocery run", null));
+    expect(appendedSnapshots(onNew)).toEqual([
+      { id: "n1", name: "note", state: { tasks: [] } },
+    ]);
+  });
+
+  it("leaves a non-user append unstamped, since no reader folds those", async () => {
+    const { thread, onNew } = externalThread({
+      composerMetadata: live({ v: 1 }),
+    });
+    await thread.append({
+      ...userMessage("an assistant turn", null),
+      role: "assistant",
+    } as AppendMessage);
+    expect(appendedSnapshots(onNew)).toBeUndefined();
+  });
+
+  it("re-stamps the baseline when the edited message carried the only snapshot", async () => {
+    const messages = [snapshotMessage("msg-1", { v: 1 })];
+    const { thread, onEdit } = externalThread({
+      messages,
+      composerMetadata: live({ v: 1 }),
+    });
+    // An edit branches off the edited message's parent, so the prefix it is
+    // gated against excludes the edited message and its snapshot.
+    await thread.append({ ...userMessage("new", null), sourceId: "msg-1" });
+    expect(appendedSnapshots(onEdit)).toEqual([
+      { id: "n1", name: "note", state: { v: 1 } },
+    ]);
+  });
+
+  it("stamps nothing when the edit's branch prefix already carries the live state", async () => {
+    const messages = [
+      snapshotMessage("parent-1", { v: 1 }),
+      snapshotMessage("msg-1", { v: 1 }),
+    ];
+    const { thread, onEdit } = externalThread({
+      messages,
+      composerMetadata: live({ v: 1 }),
+    });
+    await thread.append({
+      ...userMessage("new", "parent-1"),
+      sourceId: "msg-1",
+    });
+    expect(appendedSnapshots(onEdit)).toBeUndefined();
+  });
+
+  it("stamps on the local runtime core too", async () => {
+    const run = vi.fn(async function* () {
+      yield { content: [{ type: "text" as const, text: "ok" }] };
+    } as unknown as ChatModelAdapter["run"]);
+    const core = new LocalRuntimeCore(
+      { adapters: { chatModel: { run } } },
+      undefined,
+    );
+    core.registerModelContextProvider({
+      getModelContext: () => ({
+        unstable_composerMetadata: live({ v: 1 }),
+      }),
+    });
+    const thread = core.threads.getMainThreadRuntimeCore();
+    await thread.append(userMessage("hello", null));
+    expect(thread.messages[0]!.metadata.custom?.interactables).toEqual([
+      { id: "n1", name: "note", state: { v: 1 } },
+    ]);
+  });
+});
+
+describe("queued sends gate at dispatch, not at enqueue", () => {
+  const queuedExternalThread = () => {
+    const run = vi.fn();
+    const onNew = vi.fn(async () => {});
+    const queue = createMessageQueue({ run });
+    let composerMetadata: Record<string, unknown> = live({ v: 1 });
+    const store = (messages: ThreadMessage[]) => ({
+      messages,
+      onNew,
+      queue: queue.adapter,
+    });
+    const core = new ExternalStoreRuntimeCore(store([]));
+    core.registerModelContextProvider({
+      getModelContext: () => ({ unstable_composerMetadata: composerMetadata }),
+    });
+    return {
+      thread: core.threads.getMainThreadRuntimeCore(),
+      queue,
+      run,
+      advanceRecord: (messages: ThreadMessage[], live: unknown) => {
+        composerMetadata = live as Record<string, unknown>;
+        core.setAdapter(store(messages));
+      },
+    };
+  };
+
+  const dispatched = (run: ReturnType<typeof vi.fn>) =>
+    (run.mock.calls[0]![0] as AppendMessage).metadata?.custom?.interactables;
+
+  it("drops the enqueue-time stamp once the run has told the model that state", async () => {
+    const t = queuedExternalThread();
+    t.queue.notifyBusy();
+    await t.thread.append(userMessage("queued", null));
+    expect(t.run).not.toHaveBeenCalled();
+
+    // the run moved the record to v:2 and the live state matches it
+    t.advanceRecord([snapshotMessage("msg-1", { v: 2 })], live({ v: 2 }));
+    t.queue.notifyIdle();
+
+    expect(t.run).toHaveBeenCalledTimes(1);
+    expect(dispatched(t.run)).toBeUndefined();
+  });
+
+  it("stamps the state as of the flush when it still diverges", async () => {
+    const t = queuedExternalThread();
+    t.queue.notifyBusy();
+    await t.thread.append(userMessage("queued", null));
+
+    // the run moved the record to v:2 while the user edited on to v:3
+    t.advanceRecord([snapshotMessage("msg-1", { v: 2 })], live({ v: 3 }));
+    t.queue.notifyIdle();
+
+    expect(dispatched(t.run)).toEqual([
+      { id: "n1", name: "note", state: { v: 3 } },
+    ]);
+  });
+
+  // interrupt dispatches synchronously inside `append`, so nothing can go
+  // stale in between; what matters is that this second dispatch point stamps
+  // at all, the way advance does.
+  it("stamps on the interrupt path a mid-run steer takes", async () => {
+    const run = vi.fn();
+    const cancel = vi.fn();
+    const onNew = vi.fn(async () => {});
+    const queue = createMessageQueue({ run, cancel });
+    let composerMetadata: Record<string, unknown> = live({ v: 1 });
+    const store = (messages: ThreadMessage[]) => ({
+      messages,
+      onNew,
+      queue: queue.adapter,
+      isRunning: true,
+    });
+    const core = new ExternalStoreRuntimeCore(store([]));
+    core.registerModelContextProvider({
+      getModelContext: () => ({ unstable_composerMetadata: composerMetadata }),
+    });
+    const thread = core.threads.getMainThreadRuntimeCore();
+
+    queue.notifyBusy();
+    composerMetadata = live({ v: 2 });
+    // isRunning routes this to the steer lane, and a cancel-capable driver
+    // makes that an interrupt rather than a buffered enqueue
+    await thread.append(
+      userMessage("steered", thread.messages.at(-1)?.id ?? null),
+    );
+
+    expect(cancel).toHaveBeenCalled();
+    expect(dispatched(run)).toEqual([
+      { id: "n1", name: "note", state: { v: 2 } },
+    ]);
+  });
+
+  it("re-points a flushed message at the tail it was gated against", async () => {
+    const t = queuedExternalThread();
+    t.queue.notifyBusy();
+    await t.thread.append(userMessage("queued", null));
+
+    t.advanceRecord([snapshotMessage("msg-1", { v: 2 })], live({ v: 3 }));
+    t.queue.notifyIdle();
+
+    // routing and gating have to name the same point, or the branch the
+    // message lands on is not the one its snapshot was computed against
+    expect((t.run.mock.calls[0]![0] as AppendMessage).parentId).toBe("msg-1");
+  });
+
+  it("leaves a hand-rolled adapter stamping at enqueue", async () => {
+    const enqueue = vi.fn();
+    const onNew = vi.fn(async () => {});
+    const core = new ExternalStoreRuntimeCore({
+      messages: [],
+      onNew,
+      // no __internal_setDispatchTransform: nothing re-gates on the way out,
+      // so the runtime has to stamp before handing the message over
+      queue: {
+        items: [],
+        steerItems: [],
+        enqueue,
+        steer: vi.fn(),
+        move: vi.fn(),
+        edit: vi.fn(),
+        remove: vi.fn(),
+      },
+    });
+    core.registerModelContextProvider({
+      getModelContext: () => ({ unstable_composerMetadata: live({ v: 1 }) }),
+    });
+    const thread = core.threads.getMainThreadRuntimeCore();
+
+    await thread.append(userMessage("queued", null));
+    expect(dispatched(enqueue)).toEqual([
+      { id: "n1", name: "note", state: { v: 1 } },
+    ]);
+  });
+
+  it("gates a queued local send against the tail it flushes onto", async () => {
+    let composerMetadata: Record<string, unknown> = live({ v: 1 });
+    let releaseFirstRun!: () => void;
+    const firstRunGate = new Promise<void>((resolve) => {
+      releaseFirstRun = resolve;
+    });
+    let runs = 0;
+    const chatModel: ChatModelAdapter = {
+      run: async function* () {
+        if (runs++ === 0) await firstRunGate;
+        yield { content: [{ type: "text" as const, text: "ok" }] };
+      },
+    };
+    const core = new LocalRuntimeCore(
+      { adapters: { chatModel }, unstable_enableMessageQueue: true },
+      undefined,
+    );
+    core.registerModelContextProvider({
+      getModelContext: () => ({ unstable_composerMetadata: composerMetadata }),
+    });
+    const thread = core.threads.getMainThreadRuntimeCore();
+    const settle = () => new Promise((resolve) => setTimeout(resolve, 20));
+
+    void thread.append(userMessage("first", null));
+    await settle();
+
+    // the first run is still open, so this one buffers with the live state at v:1
+    void thread.append(
+      userMessage("second", thread.messages.at(-1)?.id ?? null),
+    );
+    await settle();
+    expect(
+      thread.messages.some(
+        (m) => m.role === "user" && getThreadMessageText(m) === "second",
+      ),
+    ).toBe(false);
+
+    composerMetadata = live({ v: 2 });
+    releaseFirstRun();
+    await settle();
+    await settle();
+
+    const second = thread.messages.find(
+      (m) => m.role === "user" && getThreadMessageText(m) === "second",
+    );
+    expect(second!.metadata.custom?.interactables).toEqual([
+      { id: "n1", name: "note", state: { v: 2 } },
+    ]);
+  });
+});

package/src/tests/base-composer-runtime-core-send.test.ts

@@ -3,6 +3,8 @@ import { DefaultThreadComposerRuntimeCore } from "../runtime/base/default-thread
 import type { AttachmentAdapter } from "../adapters/attachment";
 import type { ThreadRuntimeCore } from "../runtime/interfaces/thread-runtime-core";
 import type { PendingAttachment } from "../types/attachment";
+import { MessageNotSentError } from "../types/error";
+import { BaseComposerRuntimeCore } from "../runtime/base/base-composer-runtime-core";
 
 const makeAdapter = (
   overrides: Partial<AttachmentAdapter> = {},
@@ -571,3 +573,127 @@ describe("BaseComposerRuntimeCore send event listener isolation", () => {
     });
   });
 });
+
+describe("BaseComposerRuntimeCore.send restore-on-undispatched", () => {
+  const rejectableAppend = () => {
+    let reject!: (error: unknown) => void;
+    const append = vi.fn(
+      () =>
+        new Promise<void>((_resolve, rejectAppend) => {
+          reject = rejectAppend;
+        }),
+    );
+    return { append, reject: (error: unknown) => reject(error) };
+  };
+
+  it("restores text, quote, and attachments when the message was never sent", async () => {
+    const { append, reject } = rejectableAppend();
+    const { composer } = makeComposer(makeAdapter(), append);
+
+    composer.setText("hello");
+    await composer.addAttachment(textFile());
+    composer.setQuote({ text: "quoted", messageId: "m-1" });
+
+    await composer.send();
+    expect(composer.text).toBe("");
+    expect(composer.attachments).toHaveLength(0);
+
+    reject(new MessageNotSentError());
+
+    await vi.waitFor(() => expect(composer.text).toBe("hello"));
+    expect(composer.quote).toEqual({ text: "quoted", messageId: "m-1" });
+    expect(composer.attachments).toHaveLength(1);
+    expect(composer.attachments[0]!.status).toEqual({ type: "complete" });
+    expect(append).toHaveBeenCalledTimes(1);
+  });
+
+  it("leaves the composer empty for any other rejection", async () => {
+    const { append, reject } = rejectableAppend();
+    const { composer } = makeComposer(makeAdapter(), append);
+
+    composer.setText("hello");
+    await composer.send();
+    reject(new Error("boom"));
+
+    await vi.waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    expect(composer.text).toBe("");
+  });
+
+  it("does not clobber a draft written after the send", async () => {
+    const { append, reject } = rejectableAppend();
+    const { composer } = makeComposer(makeAdapter(), append);
+
+    composer.setText("hello");
+    await composer.send();
+    composer.setText("new draft");
+    reject(new MessageNotSentError());
+
+    await vi.waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    expect(composer.text).toBe("new draft");
+  });
+
+  it("does not restore a draft a reset discarded", async () => {
+    const { append, reject } = rejectableAppend();
+    const { composer } = makeComposer(makeAdapter(), append);
+
+    composer.setText("hello");
+    await composer.send();
+    await composer.reset();
+    reject(new MessageNotSentError());
+
+    await vi.waitFor(() => expect(append).toHaveBeenCalledTimes(1));
+    expect(composer.text).toBe("");
+  });
+});
+
+describe("BaseComposerRuntimeCore.send restore-on-undispatched edge cases", () => {
+  class SyncThrowComposerCore extends BaseComposerRuntimeCore {
+    public readonly error = new MessageNotSentError();
+    public get canCancel() {
+      return false;
+    }
+    public get canSend() {
+      return true;
+    }
+    protected getAttachmentAdapter() {
+      return undefined;
+    }
+    protected getDictationAdapter() {
+      return undefined;
+    }
+    protected handleSend(): void {
+      throw this.error;
+    }
+    protected handleCancel(): void {}
+  }
+
+  it("restores the draft when handleSend throws synchronously", async () => {
+    const composer = new SyncThrowComposerCore();
+    composer.setText("hello");
+
+    await expect(composer.send()).rejects.toThrow(composer.error);
+
+    expect(composer.text).toBe("hello");
+  });
+
+  it("restores only the most recent of several queued drafts", async () => {
+    const rejects: ((error: unknown) => void)[] = [];
+    const append = vi.fn(
+      () =>
+        new Promise<void>((_resolve, reject) => {
+          rejects.push(reject);
+        }),
+    );
+    const { composer } = makeComposer(makeAdapter(), append);
+
+    composer.setText("one");
+    await composer.send();
+    composer.setText("two");
+    await composer.send();
+
+    rejects[0]!(new MessageNotSentError());
+    rejects[1]!(new MessageNotSentError());
+
+    await vi.waitFor(() => expect(composer.text).toBe("two"));
+  });
+});

package/src/tests/base-composer-runtime-core.test.ts

@@ -14,8 +14,14 @@ class TestComposerCore extends BaseComposerRuntimeCore {
   protected getAttachmentAdapter() {
     return this._attachmentAdapter;
   }
+  private _dictationAdapter: DictationAdapter | undefined;
+
   protected getDictationAdapter(): DictationAdapter | undefined {
-    return undefined;
+    return this._dictationAdapter;
+  }
+
+  setDictationAdapter(adapter: DictationAdapter | undefined) {
+    this._dictationAdapter = adapter;
   }
 
   setAttachmentAdapter(adapter: AttachmentAdapter | undefined) {
@@ -420,3 +426,104 @@ describe("BaseComposerRuntimeCore", () => {
     });
   });
 });
+
+describe("BaseComposerRuntimeCore.restoreDraft", () => {
+  it("refuses while the composer holds something of its own", () => {
+    const composer = new TestComposerCore();
+    composer.setText("mine");
+
+    expect(composer.restoreDraft({ text: "returned" })).toBe(false);
+    expect(composer.text).toBe("mine");
+  });
+
+  it("rebases a live dictation session onto the restored text", () => {
+    let emitSpeech!: (result: { transcript: string }) => void;
+    const composer = new TestComposerCore();
+    composer.setDictationAdapter({
+      listen: () => ({
+        status: { type: "running" },
+        stop: async () => {},
+        cancel: () => {},
+        onSpeechStart: () => () => {},
+        onSpeechEnd: () => () => {},
+        onSpeech: (callback) => {
+          emitSpeech = callback;
+          return () => {};
+        },
+      }),
+    });
+
+    composer.startDictation();
+    expect(composer.restoreDraft({ text: "returned" })).toBe(true);
+    emitSpeech({ transcript: "and this" });
+
+    expect(composer.text).toBe("returned and this");
+  });
+});
+
+describe("BaseComposerRuntimeCore.retractDraft", () => {
+  const makeCompleteAttachment = (id: string) => ({
+    id,
+    type: "file" as const,
+    name: "spec.pdf",
+    contentType: "application/pdf",
+    status: { type: "complete" as const },
+    content: [],
+  });
+
+  it("clears the composer while it still holds the exact draft", () => {
+    const composer = new TestComposerCore();
+    const draft = {
+      text: "returned",
+      quote: { text: "quoted", messageId: "m-1" },
+      attachments: [makeCompleteAttachment("a1")],
+    };
+
+    expect(composer.restoreDraft(draft)).toBe(true);
+    composer.retractDraft(draft);
+
+    expect(composer.text).toBe("");
+    expect(composer.quote).toBeUndefined();
+    expect(composer.attachments).toEqual([]);
+  });
+
+  it("leaves an edited text alone", () => {
+    const composer = new TestComposerCore();
+    const draft = { text: "returned" };
+
+    expect(composer.restoreDraft(draft)).toBe(true);
+    composer.setText("edited");
+    composer.retractDraft(draft);
+
+    expect(composer.text).toBe("edited");
+  });
+
+  it("leaves a changed quote alone", () => {
+    const composer = new TestComposerCore();
+    const draft = { text: "returned", quote: { text: "quoted" } };
+
+    expect(composer.restoreDraft(draft)).toBe(true);
+    composer.setQuote({ text: "other" });
+    composer.retractDraft(draft);
+
+    expect(composer.text).toBe("returned");
+    expect(composer.quote).toEqual({ text: "other" });
+  });
+
+  it("leaves changed attachments alone", () => {
+    const composer = new TestComposerCore();
+    const draft = {
+      text: "returned",
+      attachments: [makeCompleteAttachment("a1")],
+    };
+
+    expect(composer.restoreDraft(draft)).toBe(true);
+    composer.retractDraft({
+      ...draft,
+      attachments: [makeCompleteAttachment("a1")],
+    });
+
+    expect(composer.text).toBe("returned");
+    expect(composer.attachments).toHaveLength(1);
+  });
+});

package/src/tests/default-edit-composer-runtime-core.test.ts

@@ -229,74 +229,6 @@ describe("DefaultEditComposerRuntimeCore", () => {
     });
   });
 
-  describe("interactables snapshot gating", () => {
-    const live = (state: unknown) => ({
-      interactables: [{ id: "n1", name: "note", state }],
-    });
-    const snapshotMessage = (id: string, state: unknown): ThreadMessage =>
-      makeUserMessage({
-        id,
-        metadata: {
-          custom: { interactables: [{ id: "n1", name: "note", state }] },
-        },
-      } as Partial<ThreadMessage>);
-
-    it("re-stamps the baseline when the edited message carried the only snapshot", async () => {
-      // The new branch's prefix (before the edited message) has no snapshot,
-      // so the live (unchanged) state becomes the branch's baseline.
-      const edited = snapshotMessage("msg-1", { v: 1 });
-      const { runtime, append } = makeRuntime({
-        messages: [edited],
-        composerMetadata: live({ v: 1 }),
-      });
-      const composer = new DefaultEditComposerRuntimeCore(runtime, () => {}, {
-        parentId: null,
-        message: edited,
-      });
-      composer.setText("new");
-      await composer.send();
-      const appended = append.mock.calls[0]![0] as AppendMessage;
-      expect(appended.metadata?.custom?.interactables).toEqual([
-        { id: "n1", name: "note", state: { v: 1 } },
-      ]);
-    });
-
-    it("stamps the newest live state when the interactable was edited after the original message", async () => {
-      const edited = snapshotMessage("msg-1", { v: 1 });
-      const { runtime, append } = makeRuntime({
-        messages: [edited],
-        composerMetadata: live({ v: 2 }),
-      });
-      const composer = new DefaultEditComposerRuntimeCore(runtime, () => {}, {
-        parentId: null,
-        message: edited,
-      });
-      composer.setText("new");
-      await composer.send();
-      const appended = append.mock.calls[0]![0] as AppendMessage;
-      expect(appended.metadata?.custom?.interactables).toEqual([
-        { id: "n1", name: "note", state: { v: 2 } },
-      ]);
-    });
-
-    it("stamps nothing when the branch prefix already carries the live state", async () => {
-      const parent = snapshotMessage("parent-1", { v: 1 });
-      const edited = makeUserMessage({ id: "msg-1" });
-      const { runtime, append } = makeRuntime({
-        messages: [parent, edited],
-        composerMetadata: live({ v: 1 }),
-      });
-      const composer = new DefaultEditComposerRuntimeCore(runtime, () => {}, {
-        parentId: "parent-1",
-        message: edited,
-      });
-      composer.setText("new");
-      await composer.send();
-      const appended = append.mock.calls[0]![0] as AppendMessage;
-      expect(appended.metadata?.custom?.interactables).toBeUndefined();
-    });
-  });
-
   describe("non-user messages", () => {
     it("does not lift non-text parts for assistant messages", () => {
       const { runtime } = makeRuntime();