@assistant-ui/core 0.3.17 → 0.3.19

package/src/store/scopes/thread.ts

@@ -18,6 +18,7 @@ import type { ModelContext } from "../../model-context/types";
 import type { MessageMethods, MessageState } from "./message";
 import type { ComposerMethods, ComposerState } from "./composer";
 import type { SuggestionsMethods } from "./suggestions";
+import type { TaskMethods, TaskState } from "./task";
 
 export type ThreadState = {
   /**
@@ -44,6 +45,10 @@ export type ThreadState = {
    * The messages in the currently selected branch of the thread.
    */
   readonly messages: readonly MessageState[];
+  /**
+   * Child work derived from the thread's tool calls: every tool call that carries a nested conversation, in document order with nested tasks after their parent. The array keeps its identity while no task changed.
+   */
+  readonly tasks: readonly TaskState[];
   /**
    * The thread state.
    * @deprecated This feature is experimental
@@ -76,6 +81,10 @@ export type ThreadMethods = {
    * The suggestions shown for this thread.
    */
   suggestions(): SuggestionsMethods;
+  /**
+   * Access a task by index or toolCallId; an id resolves the first task with that toolCallId in document order.
+   */
+  task(selector: { index: number } | { id: string }): TaskMethods;
   /**
    * Append a new message to the thread.
    *
@@ -144,6 +153,13 @@ export type ThreadMeta = {
 };
 
 export type ThreadEvents = {
+  "thread.toolApprovalAnswered": {
+    threadId: string;
+    messageId: string;
+    toolCallId: string;
+    toolName: string;
+    approved: boolean;
+  };
   /**
    * A run started on this thread. Also observable as `isRunning` flipping to
    * `true` in thread state.
@@ -155,6 +171,10 @@ export type ThreadEvents = {
    * state.
    */
   "thread.runEnd": { threadId: string };
+  /** The user stopped the run in progress on this thread. */
+  "thread.cancelRun": { threadId: string };
+  /** The user started a voice session on this thread. */
+  "thread.voiceStarted": { threadId: string };
   /**
    * The thread transitioned from new to initialized. Fires before the first
    * message is added, so read thread state via `useAuiState` rather than

package/src/store/scopes/threads.ts

@@ -9,6 +9,7 @@ export type ThreadsState = {
   readonly mainThreadId: string;
   readonly newThreadId: string | null;
   readonly isLoading: boolean;
+  readonly loadError: unknown;
   readonly isLoadingMore: boolean;
   readonly hasMore: boolean;
   readonly threadIds: readonly string[];

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,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", () => {
@@ -240,6 +236,80 @@ describe("MessageRepository", () => {
       expect(messages2.map((m) => m.id)).toEqual(["parent-id", "branch2-id"]);
     });
 
+    it("keeps a nested branch selected when revisiting its ancestor", () => {
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "r" }));
+      repository.addOrUpdateMessage("r", createTestMessage({ id: "a" }));
+      repository.addOrUpdateMessage("a", createTestMessage({ id: "x" }));
+      repository.addOrUpdateMessage("r", createTestMessage({ id: "b" }));
+      repository.addOrUpdateMessage("b", createTestMessage({ id: "y" }));
+
+      repository.switchToBranch("y");
+      expect(repository.getMessages().map((m) => m.id)).toEqual([
+        "r",
+        "b",
+        "y",
+      ]);
+
+      repository.switchToBranch("r");
+      expect(repository.getMessages().map((m) => m.id)).toEqual([
+        "r",
+        "b",
+        "y",
+      ]);
+
+      repository.switchToBranch("x");
+      repository.switchToBranch("r");
+      expect(repository.getMessages().map((m) => m.id)).toEqual([
+        "r",
+        "a",
+        "x",
+      ]);
+    });
+
+    it("keeps the selected root branch when deletion selects from the root", () => {
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "a" }));
+      repository.addOrUpdateMessage("a", createTestMessage({ id: "x" }));
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "b" }));
+      repository.addOrUpdateMessage("b", createTestMessage({ id: "y" }));
+
+      repository.switchToBranch("y");
+      repository.deleteMessage("y", null);
+
+      expect(repository.getMessages().map((m) => m.id)).toEqual(["b"]);
+    });
+
+    it("keeps the selected root branch after relinking a message", () => {
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "a" }));
+      repository.addOrUpdateMessage("a", createTestMessage({ id: "m" }));
+      repository.addOrUpdateMessage("m", createTestMessage({ id: "x" }));
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "b" }));
+      repository.addOrUpdateMessage("b", createTestMessage({ id: "m" }));
+
+      expect(repository.getMessages().map((m) => m.id)).toEqual([
+        "b",
+        "m",
+        "x",
+      ]);
+
+      repository.deleteMessage("x", null);
+
+      expect(repository.getMessages().map((m) => m.id)).toEqual(["b", "m"]);
+    });
+
+    it("keeps the selected branch when relinking an off-branch message", () => {
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "a" }));
+      repository.addOrUpdateMessage("a", createTestMessage({ id: "m" }));
+      repository.addOrUpdateMessage("a", createTestMessage({ id: "n" }));
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "b" }));
+      repository.addOrUpdateMessage("b", createTestMessage({ id: "n" }));
+
+      expect(repository.getMessages().map((m) => m.id)).toEqual(["a", "m"]);
+
+      repository.deleteMessage("m", null);
+
+      expect(repository.getMessages().map((m) => m.id)).toEqual(["a"]);
+    });
+
     it("should operate on a long message history without overflowing the stack", () => {
       const messageCount = 30_000;
       const messages = createLongBranchMessages(messageCount);
@@ -486,6 +556,31 @@ describe("MessageRepository", () => {
       ).toBe("parent-id");
     });
 
+    it("round-trips after reparenting under a later-added ancestor", () => {
+      const messageA = createTestMessage({ id: "A" });
+      repository.addOrUpdateMessage(null, messageA);
+      repository.addOrUpdateMessage("A", createTestMessage({ id: "B" }));
+      repository.addOrUpdateMessage("A", createTestMessage({ id: "C" }));
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "X" }));
+      repository.addOrUpdateMessage("X", messageA);
+
+      const exported = repository.export();
+      expect(exported.messages.map((m) => m.message.id)).toEqual([
+        "X",
+        "A",
+        "B",
+        "C",
+      ]);
+
+      const restored = new MessageRepository();
+      restored.import(exported);
+
+      expect(restored.headId).toBe("B");
+      expect(restored.getMessages().map((m) => m.id)).toEqual(["X", "A", "B"]);
+      restored.switchToBranch("C");
+      expect(restored.getMessages().map((m) => m.id)).toEqual(["X", "A", "C"]);
+    });
+
     it("should import repository state", () => {
       const parent = createTestMessage({ id: "parent-id" });
       const child = createTestMessage({ id: "child-id" });
@@ -941,5 +1036,22 @@ describe("MessageRepository", () => {
 
       expect(repository.headId).toBe("C");
     });
+
+    it("should advance the head when re-parenting onto the current head", () => {
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "a" }));
+      repository.addOrUpdateMessage("a", createTestMessage({ id: "h" }));
+      repository.addOrUpdateMessage(null, createTestMessage({ id: "b" }));
+      repository.addOrUpdateMessage("b", createTestMessage({ id: "c" }));
+
+      repository.addOrUpdateMessage("h", createTestMessage({ id: "c" }));
+
+      expect(repository.getMessages().map((m) => m.id)).toEqual([
+        "a",
+        "h",
+        "c",
+      ]);
+      expect(repository.headId).toBe("c");
+      expect(repository.getMessage("c").parentId).toBe("h");
+    });
   });
 });

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

@@ -22,6 +22,7 @@ import { deferred } from "./remote-thread-list-test-helpers";
 
 const EMPTY_STATE: RemoteThreadState = {
   isLoading: false,
+  loadError: undefined,
   isLoadingMore: false,
   cursor: undefined,
   newThreadId: undefined,