@assistant-ui/core 0.3.18 → 0.3.20

package/src/store/thread-tasks.test.tsx

@@ -0,0 +1,209 @@
+// @vitest-environment jsdom
+
+import { cleanup, render, waitFor } from "@testing-library/react";
+import { useAui, useAuiState } from "@assistant-ui/store";
+import { afterEach, describe, expect, it } from "vitest";
+import type { ThreadMessage } from "../types/message";
+import { AssistantRuntimeProvider } from "../react/AssistantRuntimeProvider";
+import { useExternalStoreRuntime } from "../react/runtimes/useExternalStoreRuntime";
+import type { TaskState } from "./scopes/task";
+
+const taskMessage = (
+  nestedStatus: { type: "running" } | { type: "complete"; reason: "stop" },
+  result?: unknown,
+) =>
+  ({
+    id: "message-1",
+    role: "assistant",
+    createdAt: new Date(0),
+    content: [
+      {
+        type: "tool-call",
+        toolCallId: "delegate-1",
+        toolName: "delegate",
+        args: { topic: "research" },
+        argsText: '{"topic":"research"}',
+        ...(result !== undefined ? { result } : undefined),
+        messages: [
+          {
+            id: "nested-message-1",
+            role: "assistant",
+            createdAt: new Date(0),
+            content: [],
+            status: nestedStatus,
+            metadata: {
+              unstable_state: {},
+              unstable_annotations: [],
+              unstable_data: [],
+              steps: [],
+              custom: {},
+            },
+          },
+        ],
+      },
+    ],
+    status: nestedStatus,
+    metadata: {
+      unstable_state: {},
+      unstable_annotations: [],
+      unstable_data: [],
+      steps: [],
+      custom: {},
+    },
+  }) as ThreadMessage;
+
+afterEach(cleanup);
+
+describe("thread tasks", () => {
+  it("exposes task state and accessors through an external-store runtime", async () => {
+    const captured: {
+      tasks?: readonly TaskState[];
+      byId?: TaskState;
+      byIndex?: TaskState;
+    } = {};
+
+    const Probe = () => {
+      const tasks = useAuiState((s) => s.thread.tasks);
+      const aui = useAui();
+      captured.tasks = tasks;
+      if (tasks.length > 0) {
+        captured.byId = aui.thread.task({ id: "delegate-1" }).getState();
+        captured.byIndex = aui.thread.task({ index: 0 }).getState();
+      }
+      return null;
+    };
+
+    const App = ({ messages }: { messages: readonly ThreadMessage[] }) => {
+      const runtime = useExternalStoreRuntime({
+        messages,
+        convertMessage: (message) => message,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Probe />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    const view = render(<App messages={[taskMessage({ type: "running" })]} />);
+
+    await waitFor(() => expect(captured.tasks).toHaveLength(1));
+    expect(captured.tasks?.[0]).toMatchObject({
+      status: { type: "running" },
+      toolName: "delegate",
+    });
+    expect(captured.byId).toMatchObject({ messageId: "message-1" });
+    expect(captured.byIndex).toBe(captured.byId);
+
+    view.rerender(
+      <App
+        messages={[taskMessage({ type: "complete", reason: "stop" }, "done")]}
+      />,
+    );
+
+    await waitFor(() =>
+      expect(captured.tasks?.[0]?.status).toEqual({ type: "complete" }),
+    );
+    expect(captured.tasks?.[0]).toMatchObject({
+      toolName: "delegate",
+      messageId: "message-1",
+      result: "done",
+    });
+    expect(captured.tasks?.[0]?.messages).toHaveLength(1);
+  });
+
+  it("survives nested payloads that repeat ids and resolves the first task in document order", async () => {
+    const nested = (id: string, toolCallId: string) =>
+      ({
+        id,
+        role: "assistant",
+        createdAt: new Date(0),
+        content: [
+          {
+            type: "tool-call",
+            toolCallId,
+            toolName: "search",
+            args: {},
+            argsText: "{}",
+            result: "ok",
+            messages: [],
+          },
+        ],
+        status: { type: "complete", reason: "stop" },
+        metadata: {
+          unstable_state: {},
+          unstable_annotations: [],
+          unstable_data: [],
+          steps: [],
+          custom: {},
+        },
+      }) as ThreadMessage;
+    const message = {
+      id: "message-1",
+      role: "assistant",
+      createdAt: new Date(0),
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "delegate-a",
+          toolName: "delegate",
+          args: {},
+          argsText: "{}",
+          result: "ok",
+          messages: [nested("nested", "call_1")],
+        },
+        {
+          type: "tool-call",
+          toolCallId: "delegate-b",
+          toolName: "delegate",
+          args: {},
+          argsText: "{}",
+          result: "ok",
+          messages: [nested("nested", "call_1")],
+        },
+      ],
+      status: { type: "complete", reason: "stop" },
+      metadata: {
+        unstable_state: {},
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    } as ThreadMessage;
+    const captured: { ids?: readonly string[]; first?: TaskState } = {};
+    const Probe = () => {
+      const tasks = useAuiState((s) => s.thread.tasks);
+      const aui = useAui();
+      captured.ids = tasks.map((task) => task.id);
+      if (tasks.length === 4) {
+        captured.first = aui.thread.task({ id: "call_1" }).getState();
+      }
+      return null;
+    };
+    const App = () => {
+      const runtime = useExternalStoreRuntime({
+        messages: [message],
+        convertMessage: (m) => m,
+        onNew: async () => {},
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Probe />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    render(<App />);
+
+    await waitFor(() => expect(captured.ids).toHaveLength(4));
+    expect(captured.ids).toEqual([
+      "delegate-a",
+      "call_1",
+      "delegate-b",
+      "call_1",
+    ]);
+    expect(captured.first).toMatchObject({ parentTaskId: "delegate-a" });
+  });
+});

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,84 @@ 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: { subscribe: (callback: () => void) => () => void } = 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: {
+      unstable_on: (
+        event: string,
+        callback: (payload?: unknown) => void,
+      ) => () => void;
+    } = 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/attachment-adapters.test.ts

@@ -63,7 +63,22 @@ describe("SimpleTextAttachmentAdapter", () => {
     expect(complete.content).toEqual([
       {
         type: "text",
-        text: "<attachment name=notes.md>\n- retry with backoff\n- cap attempts at 3\n</attachment>",
+        text: '<attachment name="notes.md">\n- retry with backoff\n- cap attempts at 3\n</attachment>',
+      },
+    ]);
+  });
+
+  it("quotes and escapes filenames in the attachment wrapper", async () => {
+    const adapter = new SimpleTextAttachmentAdapter();
+    const pending = (await adapter.add({
+      file: makeFile("hello", `my "notes" & <draft>'s.md`),
+    })) as PendingAttachment;
+    const complete = await adapter.send(pending);
+
+    expect(complete.content).toEqual([
+      {
+        type: "text",
+        text: '<attachment name="my &quot;notes&quot; &amp; &lt;draft&gt;\'s.md">\nhello\n</attachment>',
       },
     ]);
   });