@assistant-ui/core 0.3.20 → 0.3.21

package/src/store/runtime-clients/thread-runtime-client.test.tsx

@@ -7,7 +7,7 @@ import {
   type AssistantClient,
 } from "@assistant-ui/store";
 import { afterEach, describe, expect, it, vi } from "vitest";
-import type { ThreadListItemState } from "../../runtime/api/bindings";
+import type { ThreadListItemRuntimeState } from "../../runtime/api/bindings";
 import { ThreadRuntimeImpl } from "../../runtime/api/thread-runtime";
 import type { ExternalStoreAdapter } from "../../runtimes/external-store/external-store-adapter";
 import { ExternalStoreThreadRuntimeCore } from "../../runtimes/external-store/external-store-thread-runtime-core";
@@ -19,7 +19,7 @@ const path = {
   threadSelector: { type: "main" as const },
 };
 
-const threadListItem: ThreadListItemState = {
+const threadListItem: ThreadListItemRuntimeState = {
   id: "thread-1",
   remoteId: undefined,
   externalId: undefined,
@@ -37,7 +37,10 @@ const userMessage = {
   metadata: { custom: {} },
 } as ThreadMessage;
 
-const renderThreadClient = (core: ExternalStoreThreadRuntimeCore) => {
+const renderThreadClient = (
+  core: ExternalStoreThreadRuntimeCore,
+  configureRuntime?: (runtime: ThreadRuntimeImpl) => void,
+) => {
   const runtime = new ThreadRuntimeImpl(
     {
       path,
@@ -51,6 +54,7 @@ const renderThreadClient = (core: ExternalStoreThreadRuntimeCore) => {
       subscribe: () => () => {},
     },
   );
+  configureRuntime?.(runtime);
   const captured: { current: AssistantClient | null } = { current: null };
   const App = () => (
     <AuiProvider
@@ -62,11 +66,12 @@ const renderThreadClient = (core: ExternalStoreThreadRuntimeCore) => {
       {null}
     </AuiProvider>
   );
+  let unmount!: () => void;
   act(() => {
-    render(<App />);
+    ({ unmount } = render(<App />));
   });
   if (!captured.current) throw new Error("Expected the client to mount.");
-  return { runtime, client: captured.current };
+  return { runtime, client: captured.current, unmount };
 };
 
 describe("ThreadClient", () => {
@@ -74,6 +79,76 @@ describe("ThreadClient", () => {
     cleanup();
   });
 
+  it("does not mark the runtime's last message as last while a sent message follows it", async () => {
+    const upload = Promise.withResolvers<void>();
+    const onNew = vi.fn();
+    const store: ExternalStoreAdapter<ThreadMessage> = {
+      messages: [userMessage],
+      onNew,
+      adapters: {
+        attachments: {
+          accept: "*",
+          add: async ({ file }) => ({
+            id: "f",
+            type: "file",
+            name: file.name,
+            contentType: file.type,
+            file,
+            status: { type: "requires-action", reason: "composer-send" },
+          }),
+          remove: async () => {},
+          send: async (attachment) => {
+            await upload.promise;
+            return {
+              ...attachment,
+              status: { type: "complete" },
+              content: [],
+            };
+          },
+        },
+      },
+    };
+    const core = new ExternalStoreThreadRuntimeCore(
+      { getModelContext: () => ({}) },
+      store,
+    );
+    const { runtime, client } = renderThreadClient(core);
+    const rows = () =>
+      client.thread
+        .getState()
+        .messages.map(({ id, isLast }) => ({ id, isLast }));
+
+    await act(async () => {
+      await runtime.composer.addAttachment(
+        new File(["content"], "f.txt", { type: "text/plain" }),
+      );
+      runtime.composer.setText("hello");
+    });
+    await act(async () => {
+      runtime.composer.send();
+    });
+
+    expect(rows().map(({ isLast }) => isLast)).toEqual([false, true]);
+
+    await act(async () => {
+      upload.resolve();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    expect(onNew).toHaveBeenCalledTimes(1);
+    expect(rows().map(({ isLast }) => isLast)).toEqual([false, true]);
+
+    act(() => {
+      core.__internal_setAdapter({
+        ...store,
+        messages: [userMessage, { ...userMessage, id: "u2" }],
+      });
+    });
+    expect(rows()).toEqual([
+      { id: "u1", isLast: false },
+      { id: "u2", isLast: true },
+    ]);
+  });
+
   it("renders a run cancelled before its placeholder reached the client", async () => {
     const adapter = (
       overrides: Partial<ExternalStoreAdapter<ThreadMessage>>,
@@ -107,4 +182,62 @@ describe("ThreadClient", () => {
       "u1",
     ]);
   });
+
+  it("attempts every event cleanup when one unsubscribe throws", () => {
+    const core = new ExternalStoreThreadRuntimeCore(
+      { getModelContext: () => ({}) },
+      {
+        messages: [],
+        onNew: vi.fn(),
+        onCancel: vi.fn(),
+      },
+    );
+    const cleanupError = new Error("cleanup failed");
+    const cleanupOrder: number[] = [];
+    let subscriptionCount = 0;
+    const { unmount } = renderThreadClient(core, (runtime) => {
+      vi.spyOn(runtime, "unstable_on").mockImplementation((() => {
+        const index = subscriptionCount++;
+        return () => {
+          cleanupOrder.push(index);
+          if (index === 0) throw cleanupError;
+        };
+      }) as never);
+    });
+
+    expect(() => unmount()).toThrow(cleanupError);
+    expect(subscriptionCount).toBeGreaterThan(1);
+    expect(cleanupOrder).toEqual(
+      Array.from({ length: subscriptionCount }, (_, index) => index),
+    );
+  });
+
+  it("attempts every composer event cleanup when one unsubscribe throws", () => {
+    const core = new ExternalStoreThreadRuntimeCore(
+      { getModelContext: () => ({}) },
+      {
+        messages: [],
+        onNew: vi.fn(),
+        onCancel: vi.fn(),
+      },
+    );
+    const cleanupError = new Error("cleanup failed");
+    const cleanupOrder: number[] = [];
+    let subscriptionCount = 0;
+    const { unmount } = renderThreadClient(core, (runtime) => {
+      vi.spyOn(runtime.composer, "unstable_on").mockImplementation((() => {
+        const index = subscriptionCount++;
+        return () => {
+          cleanupOrder.push(index);
+          if (index === 0) throw cleanupError;
+        };
+      }) as never);
+    });
+
+    expect(() => unmount()).toThrow(cleanupError);
+    expect(subscriptionCount).toBeGreaterThan(1);
+    expect(cleanupOrder).toEqual(
+      Array.from({ length: subscriptionCount }, (_, index) => index),
+    );
+  });
 });

package/src/store/runtime-clients/thread-runtime-client.ts

@@ -14,6 +14,8 @@ import {
   useClientResource,
 } from "@assistant-ui/store/client";
 import { ComposerClient } from "./composer-runtime-client";
+import { ThreadMessageClient } from "../clients/thread-message-client";
+import { submissionThreadMessage } from "../clients/submission-message";
 import { MessageClient } from "./message-runtime-client";
 import { ThreadSuggestions } from "../clients/suggestions";
 import {
@@ -23,17 +25,20 @@ import {
 } from "../clients/thread-tasks";
 import { useSubscribable } from "./useSubscribable";
 import type { ThreadState } from "../scopes/thread";
+import { runCleanups } from "../../subscribable/subscribable";
 
 const useMessageClientById = ({
   runtime,
   id,
   threadIdRef,
   threadId,
+  isLast,
 }: {
   runtime: ThreadRuntime;
   id: string;
   threadIdRef: RefObject<string>;
   threadId: string;
+  isLast: false | undefined;
 }) => {
   const messageRuntime = useMemo(
     () => runtime.getMessageById(id),
@@ -41,7 +46,7 @@ const useMessageClientById = ({
   );
 
   return useResource(
-    MessageClient({ runtime: messageRuntime, threadIdRef, threadId }),
+    MessageClient({ runtime: messageRuntime, threadIdRef, threadId, isLast }),
   );
 };
 
@@ -82,9 +87,7 @@ const useThreadClient = ({
       }),
     );
 
-    return () => {
-      for (const unsub of unsubscribers) unsub();
-    };
+    return () => runCleanups(unsubscribers);
   }, [runtime, emit]);
 
   const threadIdRef = useMemo(
@@ -124,20 +127,54 @@ const useThreadClient = ({
       withKey(getTaskKey(task), TaskClient({ task }), [task]),
     ),
   );
-  const messages = useClientLookup(
-    runtimeState.messages.map((m) =>
-      withKey(
+  const submission = composer.state.submission;
+  const inTransit = composer.state.inTransit;
+  // Messages the composer sent that the runtime does not show yet render after
+  // the thread's own, oldest first, so a send never leaves the conversation.
+  const pending = useMemo(
+    () => [...(inTransit ?? []), ...(submission ? [submission] : [])],
+    [inTransit, submission],
+  );
+  const pendingMessages = useMemo(
+    () => pending.map(submissionThreadMessage),
+    [pending],
+  );
+  const lastIndex = runtimeState.messages.length - 1;
+  const messages = useClientLookup([
+    ...runtimeState.messages.map((m, index) => {
+      const isLast =
+        index === lastIndex && pending.length > 0 ? false : undefined;
+      return withKey(
         m.id,
         MessageClientById({
           runtime,
           id: m.id,
           threadIdRef,
           threadId: runtimeState.threadId,
+          isLast,
         }),
-        [runtime, m.id, threadIdRef, runtimeState.threadId],
+        [runtime, m.id, threadIdRef, runtimeState.threadId, isLast],
+      );
+    }),
+    ...pending.map((row, index) =>
+      withKey(
+        row.id,
+        ThreadMessageClient({
+          message: pendingMessages[index]!,
+          submission: row,
+          index: runtimeState.messages.length + index,
+          isLast: index === pending.length - 1,
+        }),
+        [
+          pendingMessages[index],
+          row,
+          runtimeState.messages.length,
+          index,
+          pending.length,
+        ],
       ),
     ),
-  );
+  ]);
 
   const state = useMemo<ThreadState>(() => {
     return {

package/src/store/scopes/composer.ts

@@ -5,6 +5,7 @@ import type { RunConfig } from "../../types/message";
 import type { ComposerRuntime } from "../../runtime/api/composer-runtime";
 import type {
   AttachmentAddErrorReason,
+  ComposerSubmission,
   DictationState,
   SendOptions,
 } from "../../runtime/interfaces/composer-runtime-core";
@@ -36,10 +37,13 @@ export type ComposerState = {
   /**
    * Whether the composer is currently willing to send. `true` when the
    * composer is in editing mode and has non-empty content; for thread
-   * composers also requires the thread's `isSendDisabled` flag to be unset.
-   * Edit composers (saving message edits) ignore `isSendDisabled` since it
-   * is a thread-scoped gate. Cross-thread gating (running, queue capability)
-   * is layered on top by `useComposerSend`.
+   * composers also requires the thread's `isSendDisabled` flag to be unset
+   * and, while a voice session is connected, a session that takes typed text
+   * (`voice.canSendText`) and a draft without attachments. Edit composers
+   * (saving message edits) ignore `isSendDisabled` since it is a
+   * thread-scoped gate and never send while a voice session is connected.
+   * Cross-thread gating (running, queue capability) is layered on top by
+   * `useComposerSend`.
    */
   readonly canSend: boolean;
   readonly attachmentAccept: string;
@@ -63,6 +67,20 @@ export type ComposerState = {
    * Empty when no messages are queued.
    */
   readonly queue: readonly QueueItemState[];
+
+  /**
+   * The message this composer sent while its attachments are prepared. The
+   * thread renders it as its last message until the runtime takes it.
+   * Undefined when no send is being prepared.
+   */
+  readonly submission?: ComposerSubmission | undefined;
+
+  /**
+   * Messages this composer handed to the runtime that the thread does not
+   * show yet, oldest first. The thread renders them after its own messages
+   * until the runtime shows each one.
+   */
+  readonly inTransit?: readonly ComposerSubmission[] | undefined;
 };
 
 export type ComposerMethods = {

package/src/store/scopes/message.ts

@@ -2,6 +2,7 @@ import type { ThreadMessage } from "../../types/message";
 import type { RunConfig } from "../../types/message";
 import type { SpeechState } from "../../runtime/interfaces/thread-runtime-core";
 import type { MessageRuntime } from "../../runtime/api/message-runtime";
+import type { ComposerSubmission } from "../../runtime/interfaces/composer-runtime-core";
 import type { ComposerMethods, ComposerState } from "./composer";
 import type { PartMethods, PartState } from "./part";
 import type { AttachmentMethods } from "./attachment";
@@ -34,6 +35,13 @@ export type MessageState = ThreadMessage & {
   readonly isHovering: boolean;
   /** The position of this message in the thread (0 for first message) */
   readonly index: number;
+  /**
+   * Set while this row is a message the composer sent that the thread does
+   * not show as one of its own yet. Its attachments live here rather than on
+   * `attachments`, which only ever holds the ones a thread message was
+   * delivered with.
+   */
+  readonly submission?: ComposerSubmission | undefined;
 };
 
 export type MessageMethods = {

package/src/store/scopes/part.ts

@@ -4,6 +4,7 @@ import type {
   ThreadAssistantMessagePart,
   ToolApprovalResponse,
   ToolCallMessagePartStatus,
+  Unstable_ToolInteractionInput,
 } from "../../types/message";
 import type { MessagePartStatus } from "../../types/message";
 import type { MessagePartRuntime } from "../../runtime/api/message-part-runtime";
@@ -36,6 +37,13 @@ export type PartMethods = {
    * could not be recorded, so a renderer can leave its controls retryable.
    */
   respondToToolApproval(response: ToolApprovalResponse): Promise<void>;
+  /**
+   * Record a user interaction on a tool call part. Resolves once recorded and
+   * rejects when it cannot be.
+   */
+  unstable_recordInteraction?(
+    input: Unstable_ToolInteractionInput,
+  ): Promise<void>;
   __internal_getRuntime?(): MessagePartRuntime;
 };
 

package/src/subscribable/subscribable.test.ts

@@ -1,4 +1,4 @@
-import { describe, expect, it, vi } from "vitest";
+import { describe, expect, it, onTestFinished, vi } from "vitest";
 import type { SubscribableWithState } from "./subscribable";
 import {
   EventSubscriptionSubject,
@@ -114,6 +114,82 @@ describe("runCleanups", () => {
 });
 
 describe("ShallowMemoizeSubject", () => {
+  it("removes a subscriber when the initial connection throws", () => {
+    const connectionError = new Error("connection failed");
+    let state = { status: "empty" };
+    let sourceUpdate: (() => void) | undefined;
+    let shouldThrow = true;
+    const subject = new ShallowMemoizeSubject({
+      path: null,
+      getState: () => state,
+      subscribe: (callback) => {
+        if (shouldThrow) {
+          shouldThrow = false;
+          throw connectionError;
+        }
+        sourceUpdate = callback;
+        return () => {};
+      },
+    });
+    const failedSubscriber = vi.fn();
+
+    expect(() => subject.subscribe(failedSubscriber)).toThrow(connectionError);
+
+    const activeSubscriber = vi.fn();
+    subject.subscribe(activeSubscriber);
+    state = { status: "ready" };
+    sourceUpdate?.();
+
+    expect(failedSubscriber).not.toHaveBeenCalled();
+    expect(activeSubscriber).toHaveBeenCalledOnce();
+  });
+
+  it("cleans up the source when synchronization after subscribe throws", () => {
+    const synchronizationError = new Error("synchronization failed");
+    const sourceCleanup = vi.fn();
+    let readCount = 0;
+    const subject = new ShallowMemoizeSubject({
+      path: null,
+      getState: () => {
+        readCount += 1;
+        if (readCount === 2) throw synchronizationError;
+        return { status: "ready" };
+      },
+      subscribe: () => sourceCleanup,
+    });
+
+    expect(() => subject.subscribe(() => {})).toThrow(synchronizationError);
+
+    expect(sourceCleanup).toHaveBeenCalledOnce();
+  });
+
+  it("preserves a connection error when rollback cleanup also throws", () => {
+    const connectionError = new Error("synchronization failed");
+    const cleanupError = new Error("cleanup failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    onTestFinished(() => consoleError.mockRestore());
+    let readCount = 0;
+    const subject = new ShallowMemoizeSubject({
+      path: null,
+      getState: () => {
+        readCount += 1;
+        if (readCount === 2) throw connectionError;
+        return { status: "ready" };
+      },
+      subscribe: () => () => {
+        throw cleanupError;
+      },
+    });
+
+    expect(() => subject.subscribe(() => {})).toThrow(connectionError);
+    expect(consoleError).toHaveBeenCalledWith(
+      "[assistant-ui] Subscription rollback cleanup threw",
+      cleanupError,
+    );
+  });
+
   it("notifies subscribers when a state key is removed", () => {
     const source = createBinding({
       status: "running",
@@ -263,6 +339,96 @@ describe("LazyMemoizeSubject", () => {
 });
 
 describe("nested subscription swaps", () => {
+  it("cleans up the initial nested source when the outer subscribe throws", () => {
+    const connectionError = new Error("outer connection failed");
+    const innerCleanup = vi.fn();
+    const subject = new NestedSubscriptionSubject({
+      path: null,
+      getState: () => ({ subscribe: () => innerCleanup }),
+      subscribe: () => {
+        throw connectionError;
+      },
+    });
+
+    expect(() => subject.subscribe(() => {})).toThrow(connectionError);
+
+    expect(innerCleanup).toHaveBeenCalledOnce();
+  });
+
+  it("cleans up the initial event source when the outer subscribe throws", () => {
+    const connectionError = new Error("outer connection failed");
+    const innerCleanup = vi.fn();
+    const subject = new EventSubscriptionSubject({
+      event: "test",
+      binding: {
+        path: null,
+        getState: () => ({ unstable_on: () => innerCleanup }),
+        subscribe: () => {
+          throw connectionError;
+        },
+      },
+    });
+
+    expect(() => subject.subscribe(() => {})).toThrow(connectionError);
+
+    expect(innerCleanup).toHaveBeenCalledOnce();
+  });
+
+  it("preserves a nested connection error when rollback cleanup also throws", () => {
+    const connectionError = new Error("outer connection failed");
+    const cleanupError = new Error("inner cleanup failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    onTestFinished(() => consoleError.mockRestore());
+    const subject = new NestedSubscriptionSubject({
+      path: null,
+      getState: () => ({
+        subscribe: () => () => {
+          throw cleanupError;
+        },
+      }),
+      subscribe: () => {
+        throw connectionError;
+      },
+    });
+
+    expect(() => subject.subscribe(() => {})).toThrow(connectionError);
+    expect(consoleError).toHaveBeenCalledWith(
+      "[assistant-ui] Subscription rollback cleanup threw",
+      cleanupError,
+    );
+  });
+
+  it("preserves an event connection error when rollback cleanup also throws", () => {
+    const connectionError = new Error("outer connection failed");
+    const cleanupError = new Error("event cleanup failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    onTestFinished(() => consoleError.mockRestore());
+    const subject = new EventSubscriptionSubject({
+      event: "test",
+      binding: {
+        path: null,
+        getState: () => ({
+          unstable_on: () => () => {
+            throw cleanupError;
+          },
+        }),
+        subscribe: () => {
+          throw connectionError;
+        },
+      },
+    });
+
+    expect(() => subject.subscribe(() => {})).toThrow(connectionError);
+    expect(consoleError).toHaveBeenCalledWith(
+      "[assistant-ui] Subscription rollback cleanup threw",
+      cleanupError,
+    );
+  });
+
   it("connects the next nested source when the previous cleanup throws", () => {
     const cleanupError = new Error("cleanup failed");
     let outerUpdate!: () => void;
@@ -341,4 +507,35 @@ describe("nested subscription swaps", () => {
     nextEvent("payload");
     expect(listener).toHaveBeenCalledWith("payload");
   });
+
+  it.each([
+    { notifyOnRebind: true, calls: 1 },
+    { notifyOnRebind: false, calls: 0 },
+  ])(
+    "notifies on a source swap $calls time(s) with notifyOnRebind $notifyOnRebind",
+    ({ notifyOnRebind, calls }) => {
+      let outerUpdate!: () => void;
+      let current = { unstable_on: () => () => {} };
+      const subject = new EventSubscriptionSubject({
+        event: "test",
+        binding: {
+          path: null,
+          getState: () => current,
+          subscribe: (callback) => {
+            outerUpdate = callback;
+            return () => {};
+          },
+        },
+        notifyOnRebind,
+      });
+      const listener = vi.fn();
+      subject.subscribe(listener);
+
+      current = { unstable_on: () => () => {} };
+      outerUpdate();
+
+      expect(listener).toHaveBeenCalledTimes(calls);
+      if (calls > 0) expect(listener).toHaveBeenCalledWith({});
+    },
+  );
 });