@assistant-ui/core 0.3.18 → 0.3.19

package/src/subscribable/subscribable.test.ts

@@ -1,6 +1,12 @@
 import { describe, expect, it, vi } from "vitest";
 import type { SubscribableWithState } from "./subscribable";
-import { LazyMemoizeSubject, ShallowMemoizeSubject } from "./subscribable";
+import {
+  EventSubscriptionSubject,
+  LazyMemoizeSubject,
+  NestedSubscriptionSubject,
+  runCleanups,
+  ShallowMemoizeSubject,
+} from "./subscribable";
 
 type TestState = {
   status: string;
@@ -56,6 +62,57 @@ const createRebuildingBinding = (initialState: TestState) => {
   };
 };
 
+describe("runCleanups", () => {
+  it("runs every cleanup before rethrowing an error", () => {
+    const cleanupError = new Error("cleanup failed");
+    const laterCleanup = vi.fn();
+
+    expect(() =>
+      runCleanups([
+        () => {
+          throw cleanupError;
+        },
+        laterCleanup,
+      ]),
+    ).toThrow(cleanupError);
+    expect(laterCleanup).toHaveBeenCalledOnce();
+  });
+
+  it("preserves every error when multiple cleanups throw", () => {
+    const firstError = new Error("first cleanup failed");
+    const secondError = new Error("second cleanup failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      let thrown: unknown;
+      try {
+        runCleanups([
+          () => {
+            throw firstError;
+          },
+          () => {
+            throw secondError;
+          },
+        ]);
+      } catch (error) {
+        thrown = error;
+      }
+
+      expect(thrown).toBeInstanceOf(AggregateError);
+      expect((thrown as AggregateError).errors).toEqual([
+        firstError,
+        secondError,
+      ]);
+      expect(consoleError).toHaveBeenNthCalledWith(1, firstError);
+      expect(consoleError).toHaveBeenNthCalledWith(2, secondError);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+});
+
 describe("ShallowMemoizeSubject", () => {
   it("notifies subscribers when a state key is removed", () => {
     const source = createBinding({
@@ -110,6 +167,28 @@ describe("ShallowMemoizeSubject", () => {
     expect(subject.getState()).toEqual({ status: "ready" });
     expect(subscriber).not.toHaveBeenCalled();
   });
+
+  it("reconnects after the previous connection cleanup throws", () => {
+    let subscribeCount = 0;
+    const cleanupError = new Error("cleanup failed");
+    const binding: SubscribableWithState<TestState, null> = {
+      path: null,
+      getState: () => ({ status: "ready" }),
+      subscribe: () => {
+        subscribeCount += 1;
+        return () => {
+          throw cleanupError;
+        };
+      },
+    };
+    const subject = new ShallowMemoizeSubject(binding);
+
+    const unsubscribe = subject.subscribe(() => {});
+    expect(() => unsubscribe()).toThrow(cleanupError);
+
+    subject.subscribe(() => {});
+    expect(subscribeCount).toBe(2);
+  });
 });
 
 describe("LazyMemoizeSubject", () => {
@@ -182,3 +261,79 @@ describe("LazyMemoizeSubject", () => {
     expect(second).toEqual({ status: "done" });
   });
 });
+
+describe("nested subscription swaps", () => {
+  it("connects the next nested source when the previous cleanup throws", () => {
+    const cleanupError = new Error("cleanup failed");
+    let outerUpdate!: () => void;
+    let nextUpdate!: () => void;
+    const previous = {
+      subscribe: () => () => {
+        throw cleanupError;
+      },
+    };
+    const nextSubscribe = vi.fn((callback: () => void) => {
+      nextUpdate = callback;
+      return () => {};
+    });
+    const next = { subscribe: nextSubscribe };
+    let current = previous;
+    const subject = new NestedSubscriptionSubject({
+      path: null,
+      getState: () => current,
+      subscribe: (callback) => {
+        outerUpdate = callback;
+        return () => {};
+      },
+    });
+    const listener = vi.fn();
+    subject.subscribe(listener);
+
+    current = next;
+    expect(() => outerUpdate()).toThrow(cleanupError);
+    expect(nextSubscribe).toHaveBeenCalledOnce();
+    expect(listener).toHaveBeenCalledOnce();
+
+    nextUpdate();
+    expect(listener).toHaveBeenCalledTimes(2);
+  });
+
+  it("connects the next event source when the previous cleanup throws", () => {
+    const cleanupError = new Error("cleanup failed");
+    let outerUpdate!: () => void;
+    let nextEvent!: (payload?: unknown) => void;
+    const previous = {
+      unstable_on: () => () => {
+        throw cleanupError;
+      },
+    };
+    const nextSubscribe = vi.fn(
+      (_event: string, callback: (payload?: unknown) => void) => {
+        nextEvent = callback;
+        return () => {};
+      },
+    );
+    const next = { unstable_on: nextSubscribe };
+    let current = previous;
+    const subject = new EventSubscriptionSubject({
+      event: "test",
+      binding: {
+        path: null,
+        getState: () => current,
+        subscribe: (callback) => {
+          outerUpdate = callback;
+          return () => {};
+        },
+      },
+    });
+    const listener = vi.fn();
+    subject.subscribe(listener);
+
+    current = next;
+    expect(() => outerUpdate()).toThrow(cleanupError);
+    expect(nextSubscribe).toHaveBeenCalledOnce();
+
+    nextEvent("payload");
+    expect(listener).toHaveBeenCalledWith("payload");
+  });
+});

package/src/subscribable/subscribable.ts

@@ -58,6 +58,10 @@ export const notifySubscribers = <TArgs extends unknown[]>(
   }
 };
 
+export const runCleanups = (cleanups: Iterable<Unsubscribe>): void => {
+  notifySubscribers(cleanups);
+};
+
 const shallowEqualOrUndefined = <T extends object>(
   a: T | undefined,
   b: T | undefined,
@@ -136,8 +140,9 @@ export abstract class BaseSubject {
       if (this._connection) return;
       this._connection = this._connect();
     } else {
-      this._connection?.();
+      const connection = this._connection;
       this._connection = undefined;
+      connection?.();
     }
   }
 
@@ -284,17 +289,19 @@ export class NestedSubscriptionSubject<
       if (newState === lastState) return;
       lastState = newState;
 
-      innerUnsubscribe?.();
-      innerUnsubscribe = newState?.subscribe(callback);
-
-      callback();
+      const previousInner = innerUnsubscribe;
+      innerUnsubscribe = undefined;
+      try {
+        previousInner?.();
+      } finally {
+        innerUnsubscribe = newState?.subscribe(callback);
+        callback();
+      }
     };
 
     const outerUnsubscribe = this.outerSubscribe(onRuntimeUpdate);
-    return () => {
-      outerUnsubscribe?.();
-      innerUnsubscribe?.();
-    };
+    return () =>
+      runCleanups([() => outerUnsubscribe?.(), () => innerUnsubscribe?.()]);
   }
 }
 
@@ -329,14 +336,17 @@ export class EventSubscriptionSubject<
       if (newState === lastState) return;
       lastState = newState;
 
-      innerUnsubscribe?.();
-      innerUnsubscribe = newState?.unstable_on(this.config.event, callback);
+      const previousInner = innerUnsubscribe;
+      innerUnsubscribe = undefined;
+      try {
+        previousInner?.();
+      } finally {
+        innerUnsubscribe = newState?.unstable_on(this.config.event, callback);
+      }
     };
 
     const outerUnsubscribe = this.outerSubscribe(onRuntimeUpdate);
-    return () => {
-      outerUnsubscribe?.();
-      innerUnsubscribe?.();
-    };
+    return () =>
+      runCleanups([() => outerUnsubscribe?.(), () => innerUnsubscribe?.()]);
   }
 }

package/src/tests/MessageRepository.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
+import { describe, it, expect, beforeEach, vi } from "vitest";
 import {
   MessageRepository,
   ExportedMessageRepository,
@@ -59,10 +59,6 @@ describe("MessageRepository", () => {
     mockGenerateId.mockImplementation(() => `mock-id-${nextMockId++}`);
   });
 
-  afterEach(() => {
-    vi.clearAllMocks();
-  });
-
   // Core functionality tests - these test the public contract
   describe("Basic CRUD operations", () => {
     it("should add a new message to the repository", () => {

package/src/tests/auiV0Encode.test.ts

@@ -236,6 +236,36 @@ describe("auiV0Encode", () => {
     ]);
   });
 
+  it("preserves reasoning provider metadata", () => {
+    const encoded = auiV0Encode({
+      id: "m1",
+      createdAt: new Date(),
+      role: "assistant",
+      status: { type: "complete", reason: "stop" },
+      content: [
+        {
+          type: "reasoning",
+          text: "thinking",
+          providerMetadata: { "assistant-ui": { duration: 3200 } },
+        },
+      ],
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    });
+    expect(encoded.content).toEqual([
+      {
+        type: "reasoning",
+        text: "thinking",
+        providerMetadata: { "assistant-ui": { duration: 3200 } },
+      },
+    ]);
+  });
+
   it("preserves reasoning summaries and omits absent summaries", () => {
     const encoded = auiV0Encode({
       id: "m1",
@@ -639,4 +669,274 @@ describe("auiV0Decode", () => {
       { type: "data", name: "PredictState", data: '{"steps":["a","b"]}' },
     ]);
   });
+
+  it("round-trips parent IDs and tool-call state", () => {
+    const content = auiV0Encode({
+      id: "local",
+      createdAt: new Date("2026-03-15T00:00:00.000Z"),
+      role: "assistant",
+      status: { type: "complete", reason: "stop" },
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+      content: [
+        { type: "text", text: "answer", parentId: "text-parent" },
+        {
+          type: "reasoning",
+          text: "thinking",
+          parentId: "reasoning-parent",
+        },
+        {
+          type: "source",
+          sourceType: "url",
+          id: "source-url",
+          url: "https://example.com",
+          parentId: "url-parent",
+        },
+        {
+          type: "source",
+          sourceType: "document",
+          id: "source-document",
+          title: "notes",
+          mediaType: "text/plain",
+          parentId: "document-parent",
+        },
+        {
+          type: "file",
+          data: "file-1",
+          mimeType: "application/pdf",
+          parentId: "file-parent",
+        },
+        {
+          type: "tool-call",
+          toolCallId: "tool-1",
+          toolName: "review",
+          args: { document: "proposal" },
+          argsText: '{"document":"proposal"}',
+          parentId: "tool-parent",
+          interrupt: { type: "human", payload: { question: "continue?" } },
+          timing: { startedAt: 10, completedAt: 20 },
+          mcp: {
+            app: { resourceUri: "ui://review", serverId: "server-1" },
+          },
+          messages: [
+            {
+              id: "nested",
+              createdAt: new Date("2026-03-15T00:00:00.000Z"),
+              role: "assistant",
+              status: { type: "complete", reason: "stop" },
+              metadata: {
+                unstable_state: null,
+                unstable_annotations: [],
+                unstable_data: [],
+                steps: [],
+                custom: {},
+              },
+              content: [
+                { type: "text", text: "nested", parentId: "nested-parent" },
+              ],
+            },
+          ],
+        },
+        {
+          type: "generative-ui",
+          id: "ui-1",
+          parentId: "ui-parent",
+          spec: { root: { component: "Card", props: { title: "Review" } } },
+        },
+      ],
+    });
+
+    expect(content.content).toEqual([
+      { type: "text", text: "answer", parentId: "text-parent" },
+      {
+        type: "reasoning",
+        text: "thinking",
+        parentId: "reasoning-parent",
+      },
+      {
+        type: "source",
+        sourceType: "url",
+        id: "source-url",
+        url: "https://example.com",
+        parentId: "url-parent",
+      },
+      {
+        type: "source",
+        sourceType: "document",
+        id: "source-document",
+        title: "notes",
+        mediaType: "text/plain",
+        parentId: "document-parent",
+      },
+      {
+        type: "file",
+        data: "file-1",
+        mimeType: "application/pdf",
+        parentId: "file-parent",
+      },
+      {
+        type: "tool-call",
+        toolCallId: "tool-1",
+        toolName: "review",
+        args: { document: "proposal" },
+        parentId: "tool-parent",
+        interrupt: { type: "human", payload: { question: "continue?" } },
+        timing: { startedAt: 10, completedAt: 20 },
+        mcp: {
+          app: { resourceUri: "ui://review", serverId: "server-1" },
+        },
+        messages: [
+          {
+            id: "nested",
+            createdAt: "2026-03-15T00:00:00.000Z",
+            role: "assistant",
+            status: { type: "complete", reason: "stop" },
+            metadata: {
+              unstable_state: null,
+              unstable_annotations: [],
+              unstable_data: [],
+              steps: [],
+              custom: {},
+            },
+            content: [
+              { type: "text", text: "nested", parentId: "nested-parent" },
+            ],
+          },
+        ],
+      },
+      {
+        type: "generative-ui",
+        id: "ui-1",
+        parentId: "ui-parent",
+        spec: { root: { component: "Card", props: { title: "Review" } } },
+      },
+    ]);
+
+    const decoded = auiV0Decode({
+      id: "cloud",
+      parent_id: null,
+      format: "aui/v0",
+      content: content as never,
+      created_at: new Date("2026-03-15T00:00:00.000Z"),
+    });
+
+    if (decoded.message.role !== "assistant")
+      throw new Error("expected assistant");
+    expect(decoded.message.content).toEqual([
+      { type: "text", text: "answer", parentId: "text-parent" },
+      {
+        type: "reasoning",
+        text: "thinking",
+        parentId: "reasoning-parent",
+      },
+      {
+        type: "source",
+        sourceType: "url",
+        id: "source-url",
+        url: "https://example.com",
+        parentId: "url-parent",
+      },
+      {
+        type: "source",
+        sourceType: "document",
+        id: "source-document",
+        title: "notes",
+        mediaType: "text/plain",
+        parentId: "document-parent",
+      },
+      {
+        type: "file",
+        data: "file-1",
+        mimeType: "application/pdf",
+        parentId: "file-parent",
+      },
+      expect.objectContaining({
+        type: "tool-call",
+        toolCallId: "tool-1",
+        toolName: "review",
+        args: { document: "proposal" },
+        argsText: '{"document":"proposal"}',
+        parentId: "tool-parent",
+        interrupt: { type: "human", payload: { question: "continue?" } },
+        timing: { startedAt: 10, completedAt: 20 },
+        mcp: {
+          app: { resourceUri: "ui://review", serverId: "server-1" },
+        },
+        messages: [
+          expect.objectContaining({
+            id: "nested",
+            createdAt: new Date("2026-03-15T00:00:00.000Z"),
+            role: "assistant",
+            content: [
+              { type: "text", text: "nested", parentId: "nested-parent" },
+            ],
+          }),
+        ],
+      }),
+      {
+        type: "generative-ui",
+        id: "ui-1",
+        parentId: "ui-parent",
+        spec: { root: { component: "Card", props: { title: "Review" } } },
+      },
+    ]);
+  });
+
+  it("keeps audio message and attachment parts as audio", () => {
+    const content = auiV0Encode({
+      id: "local",
+      createdAt: new Date("2026-03-15T00:00:00.000Z"),
+      role: "user",
+      metadata: { custom: {} },
+      content: [
+        {
+          type: "audio",
+          audio: { data: "SUQzAw==", format: "mp3" },
+        },
+      ],
+      attachments: [
+        {
+          id: "attachment",
+          type: "file",
+          name: "recording",
+          status: { type: "complete" },
+          content: [
+            {
+              type: "audio",
+              audio: { data: "data:audio/wav;base64,UklGRg==", format: "wav" },
+            },
+          ],
+        },
+      ],
+    });
+
+    expect(content.content).toEqual([
+      { type: "audio", audio: { data: "SUQzAw==", format: "mp3" } },
+    ]);
+    expect(content.attachments?.[0]?.content).toEqual([
+      {
+        type: "audio",
+        audio: { data: "data:audio/wav;base64,UklGRg==", format: "wav" },
+      },
+    ]);
+
+    const decoded = auiV0Decode({
+      id: "cloud",
+      parent_id: null,
+      format: "aui/v0",
+      content: content as never,
+      created_at: new Date("2026-03-15T00:00:00.000Z"),
+    });
+
+    if (decoded.message.role !== "user") throw new Error("expected user");
+    expect(decoded.message.content).toEqual(content.content);
+    expect(decoded.message.attachments[0]?.content).toEqual(
+      content.attachments?.[0]?.content,
+    );
+  });
 });

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

@@ -5,6 +5,7 @@ import type { ThreadRuntimeCore } from "../runtime/interfaces/thread-runtime-cor
 import type { PendingAttachment } from "../types/attachment";
 import { MessageNotSentError } from "../types/error";
 import { BaseComposerRuntimeCore } from "../runtime/base/base-composer-runtime-core";
+import { LocalRuntimeCore } from "../runtimes/local/local-runtime-core";
 
 const makeAdapter = (
   overrides: Partial<AttachmentAdapter> = {},
@@ -379,6 +380,56 @@ describe("BaseComposerRuntimeCore.send restore-on-failure", () => {
     await removePromise;
   });
 
+  it("excludes cleared attachments when the upload finishes before adapter removal", async () => {
+    const upload = Promise.withResolvers<void>();
+    const removal = Promise.withResolvers<void>();
+    const adapter = makeAdapter({
+      send: async (attachment) => {
+        await upload.promise;
+        return { ...attachment, status: { type: "complete" }, content: [] };
+      },
+      remove: () => removal.promise,
+    });
+    const core = new LocalRuntimeCore(
+      {
+        adapters: {
+          chatModel: { run: async () => ({ content: [] }) },
+          attachments: adapter,
+        },
+      },
+      undefined,
+    );
+    const thread = core.threads.getMainThreadRuntimeCore();
+    const composer = thread.composer;
+    composer.setText("hello");
+    await composer.addAttachment(textFile());
+    await composer.addAttachment({
+      id: "ready",
+      name: "ready.txt",
+      content: [],
+    });
+
+    const sendPromise = composer.send({ startRun: false });
+    const clearPromise = composer.clearAttachments();
+    expect(composer.attachments).toEqual([]);
+    await composer.addAttachment({
+      id: "later",
+      name: "later.txt",
+      content: [],
+    });
+    upload.resolve();
+    await sendPromise;
+
+    expect(thread.messages).toMatchObject([
+      { content: [{ type: "text", text: "hello" }], attachments: [] },
+    ]);
+    expect(composer.attachments.map((attachment) => attachment.id)).toEqual([
+      "later",
+    ]);
+    removal.resolve();
+    await clearPromise;
+  });
+
   it("releases the in-flight lock on reset so a stalled send cannot brick the composer", async () => {
     const adapter = makeAdapter({
       send: () => new Promise(() => {}),