@assistant-ui/core 0.2.7 → 0.2.9

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

@@ -2,6 +2,7 @@ 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: () => ({}),
@@ -151,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"]);
   });
 });
 

package/src/tests/groupParts.test.ts

@@ -160,6 +160,76 @@ describe("groupPartByType", () => {
     expect(fn(notMcp)).toEqual(["group-tool"]);
   });
 
+  const standaloneContext = (...names: string[]) => ({
+    toolUIs: Object.fromEntries(
+      names.map((name) => [name, [{ render: () => null, standalone: true }]]),
+    ),
+  });
+
+  it("routes context-standalone tool calls through the 'standalone-tool-call' entry", () => {
+    const fn = groupPartByType({
+      "tool-call": ["group-tool"],
+      "standalone-tool-call": [],
+    });
+    const standalone = part({
+      type: "tool-call",
+      toolName: "ask_user",
+    } as Partial<PartState>);
+    const regular = part({
+      type: "tool-call",
+      toolName: "search",
+    } as Partial<PartState>);
+    expect(fn(standalone, standaloneContext("ask_user"))).toEqual([]);
+    expect(fn(regular, standaloneContext("ask_user"))).toEqual(["group-tool"]);
+    // No context → not standalone, falls through to "tool-call".
+    expect(fn(standalone)).toEqual(["group-tool"]);
+    // Registered but not standalone → also falls through to "tool-call".
+    const inlineCtx = {
+      toolUIs: { ask_user: [{ render: () => null, standalone: false }] },
+    };
+    expect(fn(standalone, inlineCtx)).toEqual(["group-tool"]);
+  });
+
+  it("routes MCP-app parts through 'standalone-tool-call' from the part alone", () => {
+    const fn = groupPartByType({
+      "tool-call": ["group-tool"],
+      "standalone-tool-call": [],
+    });
+    const mcpApp = part({
+      type: "tool-call",
+      toolName: "render",
+      mcp: { app: { resourceUri: "ui://my-app" } },
+    } as Partial<PartState>);
+    expect(fn(mcpApp)).toEqual([]);
+  });
+
+  it("routes MCP-app parts through the deprecated 'mcp-app' entry", () => {
+    const fn = groupPartByType({
+      "tool-call": ["group-tool"],
+      "mcp-app": ["group-mcp"],
+    });
+    const mcpApp = part({
+      type: "tool-call",
+      toolName: "render",
+      mcp: { app: { resourceUri: "ui://my-app" } },
+    } as Partial<PartState>);
+    expect(fn(mcpApp)).toEqual(["group-mcp"]);
+  });
+
+  it("prefers 'standalone-tool-call' over the deprecated 'mcp-app' entry", () => {
+    const fn = groupPartByType({
+      "tool-call": ["group-tool"],
+      "standalone-tool-call": ["group-standalone"],
+      "mcp-app": ["group-mcp"],
+    });
+    const mcpApp = part({
+      type: "tool-call",
+      toolName: "render",
+      mcp: { app: { resourceUri: "ui://x" } },
+    } as Partial<PartState>);
+    expect(fn(mcpApp)).toEqual(["group-standalone"]);
+  });
+
   it("tags the function with a GROUPBY_MEMO_KEY fingerprint", () => {
     const fn = groupPartByType({ reasoning: ["group-r"] });
     const memoKey = (fn as unknown as { [GROUPBY_MEMO_KEY]: string })[

package/src/tests/remote-thread-list-isLoading.test.ts

@@ -86,7 +86,6 @@ describe("RemoteThreadList isLoading lifecycle", () => {
     state.optimisticUpdate({
       execute: () => d.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -100,7 +99,6 @@ describe("RemoteThreadList isLoading lifecycle", () => {
     const promise = state.optimisticUpdate({
       execute: () => d.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -117,7 +115,6 @@ describe("RemoteThreadList isLoading lifecycle", () => {
     const promise = state.optimisticUpdate({
       execute: () => d.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -142,7 +139,6 @@ describe("RemoteThreadList isLoading lifecycle", () => {
     const promise = state.optimisticUpdate({
       execute: () => d.promise,
       loading: (s) => ({ ...s, isLoading: true }),
-      // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
       then: applyListResult,
     });
 
@@ -166,7 +162,6 @@ describe("RemoteThreadList isLoading error path", () => {
       .optimisticUpdate({
         execute: () => d.promise,
         loading: (s) => ({ ...s, isLoading: true }),
-        // biome-ignore lint/suspicious/noThenProperty: OptimisticState reducer pattern
         then: applyListResult,
       })
       .catch(() => {

package/src/types/message.ts

@@ -314,6 +314,11 @@ export type ThreadAssistantMessage = MessageCommonProps & {
     readonly steps: readonly ThreadStep[];
     readonly submittedFeedback?: { readonly type: "positive" | "negative" };
     readonly timing?: MessageTiming;
+    /**
+     * Marks a client-side optimistic placeholder. Such messages are evicted
+     * once off the head branch and are never persisted.
+     */
+    readonly isOptimistic?: boolean;
     readonly custom: Record<string, unknown>;
   };
 };
@@ -327,6 +332,7 @@ type BaseThreadMessage = {
     readonly steps?: readonly ThreadStep[];
     readonly submittedFeedback?: { readonly type: "positive" | "negative" };
     readonly timing?: MessageTiming;
+    readonly isOptimistic?: boolean;
     readonly custom: Record<string, unknown>;
   };
   readonly attachments?: ThreadUserMessage["attachments"];

package/src/utils/id.ts

@@ -5,10 +5,6 @@ export const generateId = customAlphabet(
   7,
 );
 
-const optimisticPrefix = "__optimistic__";
-export const generateOptimisticId = () => `${optimisticPrefix}${generateId()}`;
-export const isOptimisticId = (id: string) => id.startsWith(optimisticPrefix);
-
 const errorPrefix = "__error__";
 export const generateErrorMessageId = () => `${errorPrefix}${generateId()}`;
 export const isErrorMessageId = (id: string) => id.startsWith(errorPrefix);