@assistant-ui/core 0.2.6 → 0.2.8

package/src/subscribable/subscribable.ts

@@ -207,9 +207,9 @@ export class LazyMemoizeSubject<TState extends object, TPath>
 }
 
 export class NestedSubscriptionSubject<
-    TState extends Subscribable | undefined,
-    TPath,
-  >
+  TState extends Subscribable | undefined,
+  TPath,
+>
   extends BaseSubject
   implements
     SubscribableWithState<TState, TPath>,

package/src/tests/MessageRepository.test.ts

@@ -7,20 +7,14 @@ import type { ThreadMessage } from "../types/message";
 import type { TextMessagePart } from "../types/message";
 import type { ThreadMessageLike } from "../runtime/utils/thread-message-like";
 
-// Mock generateId and generateOptimisticId to make tests deterministic
+// Mock generateId to make tests deterministic
 const mockGenerateId = vi.fn();
-const mockGenerateOptimisticId = vi.fn();
-const mockIsOptimisticId = vi.fn((id: string) =>
-  id.startsWith("__optimistic__"),
-);
 
 vi.mock("../utils/id", async (importOriginal) => {
   const original = await importOriginal<typeof import("../utils/id")>();
   return {
     ...original,
     generateId: () => mockGenerateId(),
-    generateOptimisticId: () => mockGenerateOptimisticId(),
-    isOptimisticId: (id: string) => mockIsOptimisticId(id),
   };
 });
 
@@ -58,23 +52,11 @@ describe("MessageRepository", () => {
     ...overrides,
   });
 
-  /**
-   * Creates a test CoreMessage with the given overrides.
-   */
-  const createThreadMessageLike = (overrides = {}): ThreadMessageLike => ({
-    role: "assistant",
-    content: [{ type: "text", text: "Test message" }],
-    ...overrides,
-  });
-
   beforeEach(() => {
     repository = new MessageRepository();
     // Reset mocks with predictable counter-based values
     nextMockId = 1;
     mockGenerateId.mockImplementation(() => `mock-id-${nextMockId++}`);
-    mockGenerateOptimisticId.mockImplementation(
-      () => `__optimistic__mock-id-${nextMockId++}`,
-    );
   });
 
   afterEach(() => {
@@ -297,53 +279,102 @@ describe("MessageRepository", () => {
   });
 
   describe("Optimistic messages", () => {
-    it("should create an optimistic message with a unique ID", () => {
-      mockGenerateOptimisticId.mockReturnValue("__optimistic__generated-id");
+    const optimistic = (overrides = {}) =>
+      createTestMessage({
+        status: { type: "running" },
+        metadata: {
+          unstable_state: null,
+          unstable_annotations: [],
+          unstable_data: [],
+          steps: [],
+          custom: {},
+          isOptimistic: true,
+        },
+        ...overrides,
+      });
 
-      const coreMessage = createThreadMessageLike();
-      const optimisticId = repository.appendOptimisticMessage(
-        null,
-        coreMessage,
-      );
+    it("excludes optimistic messages from export()", () => {
+      const parent = createTestMessage({ id: "u" });
+      repository.addOrUpdateMessage(null, parent);
+      repository.addOrUpdateMessage("u", optimistic({ id: "placeholder" }));
+      repository.resetHead("placeholder");
 
-      expect(optimisticId).toBe("__optimistic__generated-id");
-      expect(repository.getMessage(optimisticId).message.status?.type).toBe(
-        "running",
-      );
+      const exported = repository.export();
+
+      expect(exported.messages.map((m) => m.message.id)).toEqual(["u"]);
+      // head was the optimistic placeholder; the exported head must fall back
+      // to the nearest persisted ancestor so it always resolves on import.
+      expect(exported.headId).toBe("u");
     });
 
-    it("should create an optimistic message as a child of a specified parent", () => {
-      const parent = createTestMessage({ id: "parent-id" });
+    it("round-trips through export/import without resurrecting the placeholder", () => {
+      const parent = createTestMessage({ id: "u" });
+      const real = createTestMessage({ id: "a" });
       repository.addOrUpdateMessage(null, parent);
+      repository.addOrUpdateMessage("u", real);
+      repository.resetHead("a");
 
-      const coreMessage = createThreadMessageLike();
-      const optimisticId = repository.appendOptimisticMessage(
-        "parent-id",
-        coreMessage,
-      );
+      const restored = new MessageRepository();
+      restored.import(repository.export());
 
-      const result = repository.getMessage(optimisticId);
-      expect(result.parentId).toBe("parent-id");
+      expect(restored.export().messages.map((m) => m.message.id)).toEqual([
+        "u",
+        "a",
+      ]);
     });
 
-    it("should retry generating unique optimistic IDs if initial one exists", () => {
-      mockGenerateOptimisticId.mockReturnValueOnce("__optimistic__existing-id");
+    describe("HEAD-branch invariant", () => {
+      it("evicts an off-branch optimistic sibling when resetHead moves the head", () => {
+        // u -> { client_id (optimistic), server_id (optimistic) }. When the
+        // head moves to server_id, the dangling client_id sibling is evicted.
+        const parent = createTestMessage({ id: "u" });
+        repository.addOrUpdateMessage(null, parent);
+        repository.addOrUpdateMessage("u", optimistic({ id: "client_id" }));
+        repository.addOrUpdateMessage("u", optimistic({ id: "server_id" }));
+
+        repository.resetHead("server_id");
 
-      const existingMessage = createTestMessage({
-        id: "__optimistic__existing-id",
+        expect(repository.getBranches("server_id")).toEqual(["server_id"]);
+        expect(() => repository.getMessage("client_id")).toThrow();
       });
-      repository.addOrUpdateMessage(null, existingMessage);
 
-      mockGenerateOptimisticId.mockReturnValueOnce("__optimistic__unique-id");
+      it("keeps the optimistic message that is on the head branch", () => {
+        const parent = createTestMessage({ id: "u" });
+        repository.addOrUpdateMessage(null, parent);
+        repository.addOrUpdateMessage("u", optimistic({ id: "a" }));
 
-      const coreMessage = createThreadMessageLike();
-      const optimisticId = repository.appendOptimisticMessage(
-        null,
-        coreMessage,
-      );
+        repository.resetHead("a");
+
+        expect(repository.getMessage("a").message.id).toBe("a");
+      });
 
-      expect(optimisticId).toBe("__optimistic__unique-id");
-      expect(mockGenerateOptimisticId).toHaveBeenCalledTimes(2);
+      it("never evicts real (non-optimistic) sibling branches", () => {
+        const parent = createTestMessage({ id: "u" });
+        repository.addOrUpdateMessage(null, parent);
+        repository.addOrUpdateMessage("u", createTestMessage({ id: "a1" }));
+        repository.addOrUpdateMessage("u", createTestMessage({ id: "a2" }));
+
+        repository.resetHead("a2");
+
+        // a1 is off the head branch but not optimistic, so it survives.
+        expect(repository.getBranches("a2")).toEqual(["a1", "a2"]);
+      });
+
+      it("evicts optimistic messages from the previous branch on switchToBranch", () => {
+        // Two real branches under u; the head branch (a2) carries an optimistic
+        // child. Switching to a1 must drop the optimistic message left on a2.
+        const parent = createTestMessage({ id: "u" });
+        repository.addOrUpdateMessage(null, parent);
+        repository.addOrUpdateMessage("u", createTestMessage({ id: "a1" }));
+        repository.addOrUpdateMessage("u", createTestMessage({ id: "a2" }));
+        repository.addOrUpdateMessage("a2", optimistic({ id: "opt" }));
+        repository.resetHead("opt");
+
+        repository.switchToBranch("a1");
+
+        expect(() => repository.getMessage("opt")).toThrow();
+        expect(repository.getBranches("a1")).toEqual(["a1", "a2"]);
+      });
     });
   });
 

package/src/tests/OptimisticState-delete-crash.test.ts

@@ -24,6 +24,8 @@ const createThreadState = (
   const mappingId = createThreadMappingId(threadId);
   return {
     isLoading: false,
+    isLoadingMore: false,
+    cursor: undefined,
     newThreadId: undefined,
     threadIds: status === "regular" ? [threadId] : [],
     archivedThreadIds: status === "archived" ? [threadId] : [],

package/src/tests/OptimisticState-list-race.test.ts

@@ -22,6 +22,8 @@ import { deferred } from "./remote-thread-list-test-helpers";
 
 const EMPTY_STATE: RemoteThreadState = {
   isLoading: false,
+  isLoadingMore: false,
+  cursor: undefined,
   newThreadId: undefined,
   threadIds: [],
   archivedThreadIds: [],
@@ -87,7 +89,6 @@ describe("list + delete race condition", () => {
     const listPromise = state.optimisticUpdate({
       execute: () => listDeferred.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -145,7 +146,6 @@ describe("list + delete race condition", () => {
     const listPromise = state.optimisticUpdate({
       execute: () => listDeferred.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -183,7 +183,6 @@ describe("list + delete race condition", () => {
     const listPromise = state.optimisticUpdate({
       execute: () => listDeferred.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -221,7 +220,6 @@ describe("list + delete race condition", () => {
     const listPromise = state.optimisticUpdate({
       execute: () => listDeferred.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 

package/src/tests/RemoteThreadListThreadListRuntimeCore-loadMore.test.ts

@@ -20,7 +20,7 @@ describe("RemoteThreadListThreadListRuntimeCore.loadMore", () => {
 
   it("initial list response with nextCursor sets hasMore=true", async () => {
     const adapter = makeAdapter({
-      list: vi.fn(async () => ({
+      list: vi.fn<ListFn>(async () => ({
         threads: [
           {
             status: "regular",
@@ -41,7 +41,7 @@ describe("RemoteThreadListThreadListRuntimeCore.loadMore", () => {
 
   it("absent nextCursor leaves hasMore=false", async () => {
     const adapter = makeAdapter({
-      list: vi.fn(async () => ({
+      list: vi.fn<ListFn>(async () => ({
         threads: [
           {
             status: "regular",
@@ -243,7 +243,7 @@ describe("RemoteThreadListThreadListRuntimeCore.loadMore", () => {
 
   it("__internal_setOptions clears cursor and dedup handles on adapter swap, then refetches via the new adapter", async () => {
     const firstAdapter = makeAdapter({
-      list: vi.fn(async () => ({
+      list: vi.fn<ListFn>(async () => ({
         threads: [{ status: "regular", remoteId: "old", externalId: "old" }],
         nextCursor: "old-cursor",
       })),
@@ -253,7 +253,7 @@ describe("RemoteThreadListThreadListRuntimeCore.loadMore", () => {
     expect(core.threadIds).toEqual(["old"]);
     expect(core.hasMore).toBe(true);
 
-    const secondList = vi.fn(async () => ({
+    const secondList = vi.fn<ListFn>(async () => ({
       threads: [{ status: "regular", remoteId: "new", externalId: "new" }],
     }));
     const secondAdapter = makeAdapter({ list: secondList });
@@ -290,7 +290,7 @@ describe("RemoteThreadListThreadListRuntimeCore.loadMore", () => {
     const stale = core.loadMore();
 
     const secondAdapter = makeAdapter({
-      list: vi.fn(async () => ({
+      list: vi.fn<ListFn>(async () => ({
         threads: [{ status: "regular", remoteId: "ignored", externalId: "x" }],
       })),
     });

package/src/tests/auiV0Encode.test.ts

@@ -9,7 +9,7 @@ describe("auiV0Encode", () => {
       role: "assistant",
       status: { type: "complete", reason: "stop" },
       metadata: {
-        unstable_state: undefined,
+        unstable_state: null,
         unstable_annotations: [],
         unstable_data: [],
         steps: [],

package/src/tests/composer-can-send.test.ts

@@ -102,10 +102,14 @@ describe("BaseComposerRuntimeCore.send", () => {
 describe("DefaultEditComposerRuntimeCore.canSend", () => {
   it("ignores runtime.isSendDisabled (thread-scoped flag does not block edits)", () => {
     const stub = makeRuntimeStub({ isSendDisabled: true });
-    const composer = new DefaultEditComposerRuntimeCore(stub, () => {}, {
-      parentId: null,
-      message: makeUserMessage("seed"),
-    });
+    const composer = new DefaultEditComposerRuntimeCore(
+      stub as unknown as ThreadRuntimeCore,
+      () => {},
+      {
+        parentId: null,
+        message: makeUserMessage("seed"),
+      },
+    );
 
     expect(composer.canSend).toBe(true);
   });

package/src/tests/duplicate-detection.test.ts

@@ -0,0 +1,34 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
+import { checkDuplicateCore } from "../internal/duplicate-detection";
+
+const KEY = Symbol.for("@assistant-ui/core.loaded");
+
+function reset(): void {
+  delete (globalThis as unknown as Record<symbol, unknown>)[KEY];
+}
+
+describe("checkDuplicateCore", () => {
+  let warn: ReturnType<typeof vi.spyOn>;
+
+  beforeEach(() => {
+    reset();
+    warn = vi.spyOn(console, "warn").mockImplementation(() => {});
+  });
+
+  afterEach(() => {
+    warn.mockRestore();
+    reset();
+  });
+
+  it("does not warn on the first load", () => {
+    checkDuplicateCore();
+    expect(warn).not.toHaveBeenCalled();
+  });
+
+  it("warns when a second copy registers in the same runtime", () => {
+    checkDuplicateCore();
+    checkDuplicateCore();
+    expect(warn).toHaveBeenCalledTimes(1);
+    expect(warn.mock.calls[0]![0]).toMatch(/npx assistant-ui doctor/);
+  });
+});

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

@@ -184,7 +184,7 @@ describe("ExternalStoreThreadListRuntimeCore - isMain via ThreadListRuntimeImpl"
     const core = new ExternalStoreThreadListRuntimeCore(adapter, makeFactory());
     return new ThreadListRuntimeImpl(
       core,
-      NoopThreadRuntime as unknown as Parameters<
+      NoopThreadRuntime as unknown as ConstructorParameters<
         typeof ThreadListRuntimeImpl
       >[1],
     );

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

@@ -2,18 +2,20 @@ import { describe, expect, it, vi } from "vitest";
 import { ExternalStoreThreadRuntimeCore } from "../runtimes/external-store/external-store-thread-runtime-core";
 import type { ExternalStoreAdapter } from "../runtimes/external-store/external-store-adapter";
 import type { ModelContextProvider } from "../model-context/types";
+import type { ThreadMessageLike } from "../runtime/utils/thread-message-like";
 
 const mockContextProvider: ModelContextProvider = {
   getModelContext: () => ({}),
 };
 
 const makeStore = (
-  overrides?: Partial<ExternalStoreAdapter>,
-): ExternalStoreAdapter => ({
-  messages: [],
-  onNew: vi.fn(),
-  ...overrides,
-});
+  overrides?: Partial<ExternalStoreAdapter> | Record<string, unknown>,
+): ExternalStoreAdapter =>
+  ({
+    messages: [],
+    onNew: vi.fn(),
+    ...overrides,
+  }) as ExternalStoreAdapter;
 
 describe("ExternalStoreThreadRuntimeCore - state reference stability", () => {
   describe("capabilities", () => {
@@ -150,117 +152,148 @@ describe("ExternalStoreThreadRuntimeCore - state reference stability", () => {
     expect(runtime.capabilities).toBe(capsBefore);
   });
 });
+describe("ExternalStoreThreadRuntimeCore - optimistic message reconciliation", () => {
+  type Raw = {
+    id: string;
+    role: "user" | "assistant";
+    text: string;
+    optimistic?: boolean;
+  };
+
+  const convertMessage = (m: Raw): ThreadMessageLike => ({
+    id: m.id,
+    role: m.role,
+    content: [{ type: "text", text: m.text }],
+    ...(m.optimistic && { metadata: { isOptimistic: true } }),
+  });
 
-describe("ExternalStoreThreadRuntimeCore - messages reconciliation", () => {
-  const user = { id: "u", role: "user" as const, content: [] };
-
-  it("drops ids that disappear between syncs (same length, swapped assistant id)", () => {
-    const a1 = { id: "a1", role: "assistant" as const, content: [] };
-    const a2 = { id: "a2", role: "assistant" as const, content: [] };
+  const childrenOf = (
+    runtime: ExternalStoreThreadRuntimeCore,
+    parentId: string,
+  ) =>
+    runtime
+      .export()
+      .messages.filter((m) => m.parentId === parentId)
+      .map((m) => m.message.id);
 
+  it("drops the orphaned placeholder when an optimistic id is swapped mid-run", () => {
+    const u: Raw = { id: "u", role: "user", text: "hi" };
     const runtime = new ExternalStoreThreadRuntimeCore(
       mockContextProvider,
-      makeStore({ messages: [user, a1] }),
+      makeStore({
+        messages: [
+          u,
+          { id: "client_id", role: "assistant", text: "", optimistic: true },
+        ],
+        convertMessage,
+        isRunning: true,
+      }),
     );
 
-    runtime.__internal_setAdapter(makeStore({ messages: [user, a2] }));
+    // AI SDK v6 swaps the client-generated id for the server-provided one.
+    runtime.__internal_setAdapter(
+      makeStore({
+        messages: [
+          u,
+          {
+            id: "server_id",
+            role: "assistant",
+            text: "hello",
+            optimistic: true,
+          },
+        ],
+        convertMessage,
+        isRunning: true,
+      }),
+    );
 
-    const exported = runtime.export();
-    expect(exported.messages.map((m) => m.message.id)).toEqual(["u", "a2"]);
-    const userChildren = exported.messages
-      .filter((m) => m.parentId === "u")
-      .map((m) => m.message.id);
-    expect(userChildren).toEqual(["a2"]);
+    // No phantom sibling in the live tree (what BranchPicker reads): the user
+    // message has a single child. export() omits the still-optimistic
+    // streaming message, so assert against the live getBranches instead.
+    expect(runtime.getBranches("server_id")).toEqual(["server_id"]);
   });
 
-  it("keeps prior ids when they remain in the new sync", () => {
-    const a = { id: "a", role: "assistant" as const, content: [] };
-
+  it("clears the optimistic flag once the run settles", () => {
+    const u: Raw = { id: "u", role: "user", text: "hi" };
     const runtime = new ExternalStoreThreadRuntimeCore(
       mockContextProvider,
-      makeStore({ messages: [user, a] }),
+      makeStore({
+        messages: [
+          u,
+          { id: "a", role: "assistant", text: "...", optimistic: true },
+        ],
+        convertMessage,
+        isRunning: true,
+      }),
     );
 
-    runtime.__internal_setAdapter(makeStore({ messages: [user, a] }));
+    runtime.__internal_setAdapter(
+      makeStore({
+        messages: [u, { id: "a", role: "assistant", text: "done" }],
+        convertMessage,
+        isRunning: false,
+      }),
+    );
 
-    expect(runtime.export().messages.map((m) => m.message.id)).toEqual([
-      "u",
-      "a",
-    ]);
+    const settled = runtime.export().messages.find((m) => m.message.id === "a");
+    expect(settled?.message.metadata.isOptimistic).toBeFalsy();
+    expect(childrenOf(runtime, "u")).toEqual(["a"]);
   });
 
-  it("removes trailing messages dropped from the new sync", () => {
-    const a = { id: "a", role: "assistant" as const, content: [] };
-    const u2 = { id: "u2", role: "user" as const, content: [] };
-
+  it("removes the runtime placeholder once the store provides the assistant message", () => {
+    const u: Raw = { id: "u", role: "user", text: "hi" };
     const runtime = new ExternalStoreThreadRuntimeCore(
       mockContextProvider,
-      makeStore({ messages: [user, a, u2] }),
+      makeStore({ messages: [u], convertMessage, isRunning: true }),
     );
 
-    runtime.__internal_setAdapter(makeStore({ messages: [user, a] }));
+    // Running with a trailing user message: a placeholder is appended to the
+    // live tree. export() omits optimistic messages, so inspect the live
+    // messages (which include the placeholder on the head path). It's a plain
+    // assistant message flagged optimistic (no special id scheme).
+    const live = runtime.messages;
+    expect(live).toHaveLength(2);
+    expect(live[1]!.role).toBe("assistant");
+    expect(live[1]!.metadata.isOptimistic).toBe(true);
+
+    // The store now yields the real assistant message; the placeholder (whose
+    // synthetic id never appears in the snapshot) must be gone, leaving a
+    // single child under the user message.
+    runtime.__internal_setAdapter(
+      makeStore({
+        messages: [u, { id: "a", role: "assistant", text: "done" }],
+        convertMessage,
+        isRunning: false,
+      }),
+    );
 
     expect(runtime.export().messages.map((m) => m.message.id)).toEqual([
       "u",
       "a",
     ]);
+    expect(runtime.getBranches("a")).toEqual(["a"]);
   });
 
-  it("does not crash on the next sync after cancelRun removes a leaf user", () => {
-    const userWithText = {
-      id: "u",
-      role: "user" as const,
-      content: [{ type: "text" as const, text: "hi" }],
-    };
-
+  it("keeps real sibling branches that were never flagged optimistic", () => {
+    const u: Raw = { id: "u", role: "user", text: "hi" };
     const runtime = new ExternalStoreThreadRuntimeCore(
       mockContextProvider,
       makeStore({
-        messages: [userWithText],
-        onCancel: vi.fn(),
-        isRunning: true,
-      }),
-    );
-
-    runtime.cancelRun();
-
-    expect(() => {
-      runtime.__internal_setAdapter(makeStore({ messages: [] }));
-    }).not.toThrow();
-  });
-
-  it("drops phantom sibling when convertMessage swaps the assistant id", () => {
-    type Raw = { id: string; role: "user" | "assistant"; text: string };
-    const rawU: Raw = { id: "u", role: "user", text: "hi" };
-    const rawA1: Raw = { id: "client_id", role: "assistant", text: "" };
-    const rawA2: Raw = { id: "server_id", role: "assistant", text: "" };
-
-    const convertMessage = (m: Raw) => ({
-      id: m.id,
-      role: m.role,
-      content: [{ type: "text" as const, text: m.text }],
-    });
-
-    const runtime = new ExternalStoreThreadRuntimeCore(
-      mockContextProvider,
-      makeStore({
-        messages: [rawU, rawA1] as any,
-        convertMessage: convertMessage as any,
+        messages: [u, { id: "a1", role: "assistant", text: "first" }],
+        convertMessage,
       }),
     );
 
+    // Simulates onEdit/onReload producing a new branch under the same parent;
+    // the prior branch must survive (regression guard for #4131).
     runtime.__internal_setAdapter(
       makeStore({
-        messages: [rawU, rawA2] as any,
-        convertMessage: convertMessage as any,
+        messages: [u, { id: "a2", role: "assistant", text: "second" }],
+        convertMessage,
       }),
     );
 
-    const userChildren = runtime
-      .export()
-      .messages.filter((m) => m.parentId === "u")
-      .map((m) => m.message.id);
-    expect(userChildren).toEqual(["server_id"]);
+    expect(childrenOf(runtime, "u")).toEqual(["a1", "a2"]);
   });
 });