@assistant-ui/core 0.3.17 → 0.3.18

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

@@ -1,5 +1,5 @@
 import type { Unsubscribe } from "../../types/unsubscribe";
-import { useMemo, useEffect } from "react";
+import { useMemo, useEffect, useRef } from "react";
 import { useResource, resource, withKey } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
 import { useAssistantEmit, useClientLookup } from "@assistant-ui/store/client";
@@ -59,28 +59,48 @@ const useComposerClient = ({
   threadIdRef,
   messageIdRef,
   runtime,
+  isSuggestion,
 }: {
   threadIdRef: { current: string };
   messageIdRef?: { current: string };
   runtime: ComposerRuntime;
+  isSuggestion?: ((text: string) => boolean) | undefined;
 }): ClientOutput<"composer"> => {
   const runtimeState = useSubscribable(runtime);
   const emit = useAssistantEmit();
+  const pendingSuggestion = useRef(false);
 
   // Bind composer events to event manager
   useEffect(() => {
     const unsubscribers: Unsubscribe[] = [];
 
     // Subscribe to composer events
-    for (const event of ["send", "attachmentAdd"] as const) {
-      const unsubscribe = runtime.unstable_on(event, () => {
-        emit(`composer.${event}`, {
+    const sendUnsubscribe = runtime.unstable_on("send", (payload) => {
+      const suggestion = pendingSuggestion.current;
+      pendingSuggestion.current = false;
+      emit("composer.send", {
+        threadId: threadIdRef.current,
+        ...(messageIdRef && { messageId: messageIdRef.current }),
+        chars: payload.chars,
+        attachments: payload.attachments,
+        ...(suggestion ? { suggestion: true } : undefined),
+      });
+    });
+    unsubscribers.push(sendUnsubscribe);
+
+    const attachmentUnsubscribe = runtime.unstable_on(
+      "attachmentAdd",
+      (payload) => {
+        emit("composer.attachmentAdd", {
           threadId: threadIdRef.current,
           ...(messageIdRef && { messageId: messageIdRef.current }),
+          ...(payload.contentType
+            ? { contentType: payload.contentType }
+            : undefined),
         });
-      });
-      unsubscribers.push(unsubscribe);
-    }
+      },
+    );
+    unsubscribers.push(attachmentUnsubscribe);
 
     unsubscribers.push(
       runtime.unstable_on("attachmentAddError", (payload) => {
@@ -91,6 +111,9 @@ const useComposerClient = ({
           ...(payload.attachmentId && { attachmentId: payload.attachmentId }),
           reason: payload.reason,
           message: payload.message,
+          ...(payload.contentType
+            ? { contentType: payload.contentType }
+            : undefined),
         });
       }),
     );
@@ -153,8 +176,18 @@ const useComposerClient = ({
     addAttachment: runtime.addAttachment,
     reset: runtime.reset,
     clearAttachments: runtime.clearAttachments,
-    send: runtime.send,
-    cancel: runtime.cancel,
+    send: (options) => {
+      const state = runtime.getState();
+      pendingSuggestion.current =
+        state.canSend && (isSuggestion?.(state.text) ?? false);
+      runtime.send(options);
+    },
+    cancel: () => {
+      if (!messageIdRef && runtime.getState().canCancel) {
+        emit("composer.cancel", { threadId: threadIdRef.current });
+      }
+      runtime.cancel();
+    },
     beginEdit:
       (runtime as EditComposerRuntime).beginEdit ??
       (() => {

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

@@ -1,7 +1,11 @@
-import { useMemo, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import { useResource, withKey, resource } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
-import { useClientLookup, useClientResource } from "@assistant-ui/store/client";
+import {
+  useAssistantEmit,
+  useClientLookup,
+  useClientResource,
+} from "@assistant-ui/store/client";
 import type { MessageRuntime } from "../../runtime/api/message-runtime";
 import { useSubscribable } from "./useSubscribable";
 import { liveRef } from "./liveRef";
@@ -47,11 +51,14 @@ const MessagePartByIndex = resource(useMessagePartByIndex);
 const useMessageClient = ({
   runtime,
   threadIdRef,
+  threadId,
 }: {
   runtime: MessageRuntime;
   threadIdRef: { current: string };
+  threadId: string;
 }): ClientOutput<"message"> => {
   const runtimeState = useSubscribable(runtime);
+  const emit = useAssistantEmit();
 
   const [isCopiedState, setIsCopied] = useState(false);
   const [isHoveringState, setIsHovering] = useState(false);
@@ -60,6 +67,33 @@ const useMessageClient = ({
     () => liveRef(() => runtime.getState().id),
     [runtime],
   );
+  const previousStatus = useRef(runtimeState.status);
+  const emitMessageEvent = (
+    event:
+      | "message.reload"
+      | "message.speak"
+      | "message.branchSwitched"
+      | "message.copied",
+  ) => {
+    emit(event, { threadId, messageId: runtime.getState().id });
+  };
+
+  useEffect(() => {
+    const status = runtimeState.status;
+    const previous = previousStatus.current;
+    previousStatus.current = status;
+    if (
+      status?.type === "incomplete" &&
+      status.reason === "error" &&
+      (previous?.type !== "incomplete" || previous.reason !== "error")
+    ) {
+      emit("message.error", {
+        threadId,
+        messageId: runtimeState.id,
+        reason: "error",
+      });
+    }
+  }, [runtimeState.status, runtimeState.id, emit, threadId]);
 
   const composer = useClientResource(
     ComposerClient({
@@ -114,11 +148,20 @@ const useMessageClient = ({
     composer: () => composer.methods,
 
     delete: () => runtime.delete(),
-    reload: (config) => runtime.reload(config),
-    speak: () => runtime.speak(),
+    reload: (config) => {
+      emitMessageEvent("message.reload");
+      return runtime.reload(config);
+    },
+    speak: () => {
+      emitMessageEvent("message.speak");
+      return runtime.speak();
+    },
     stopSpeaking: () => runtime.stopSpeaking(),
     submitFeedback: (feedback) => runtime.submitFeedback(feedback),
-    switchToBranch: (options) => runtime.switchToBranch(options),
+    switchToBranch: (options) => {
+      emitMessageEvent("message.branchSwitched");
+      return runtime.switchToBranch(options);
+    },
     getCopyText: () => runtime.unstable_getCopyText(),
     part: (selector) => {
       if ("index" in selector) {
@@ -136,7 +179,12 @@ const useMessageClient = ({
       }
     },
 
-    setIsCopied,
+    setIsCopied: (value) => {
+      if (value) {
+        emitMessageEvent("message.copied");
+      }
+      setIsCopied(value);
+    },
     setIsHovering,
 
     __internal_getRuntime: () => runtime,

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

@@ -52,8 +52,10 @@ const useThreadListItemClient = ({
     unarchive: () =>
       handleThreadListAction("unarchive", () => runtime.unarchive()),
     delete: () => handleThreadListAction("delete", () => runtime.delete()),
-    generateTitle: () =>
-      handleThreadListAction("generate title", () => runtime.generateTitle()),
+    generateTitle: (options) =>
+      handleThreadListAction("generate title", () =>
+        runtime.generateTitle(options),
+      ),
     initialize: runtime.initialize,
     detach: runtime.detach,
     __internal_getRuntime: () => runtime,

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

@@ -82,6 +82,7 @@ const useThreadListClient = ({
       mainThreadId: runtimeState.mainThreadId,
       newThreadId: runtimeState.newThreadId ?? null,
       isLoading: runtimeState.isLoading,
+      loadError: runtimeState.loadError,
       isLoadingMore: runtimeState.isLoadingMore,
       hasMore: runtimeState.hasMore,
       threadIds: runtimeState.threadIds,

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

@@ -1,7 +1,10 @@
 import type { Unsubscribe } from "../../types/unsubscribe";
 import type { ThreadRuntimeEventType } from "../../runtime/interfaces/thread-runtime-core";
-import type { ThreadRuntime } from "../../runtime/api/thread-runtime";
-import { useMemo, useEffect, type RefObject } from "react";
+import type {
+  CreateAppendMessage,
+  ThreadRuntime,
+} from "../../runtime/api/thread-runtime";
+import { useMemo, useEffect, useCallback, type RefObject } from "react";
 import { useResource, resource, withKey } from "@assistant-ui/tap";
 import { liveRef } from "./liveRef";
 import type { ClientOutput } from "@assistant-ui/store";
@@ -20,17 +23,21 @@ const useMessageClientById = ({
   runtime,
   id,
   threadIdRef,
+  threadId,
 }: {
   runtime: ThreadRuntime;
   id: string;
   threadIdRef: RefObject<string>;
+  threadId: string;
 }) => {
   const messageRuntime = useMemo(
     () => runtime.getMessageById(id),
     [runtime, id],
   );
 
-  return useResource(MessageClient({ runtime: messageRuntime, threadIdRef }));
+  return useResource(
+    MessageClient({ runtime: messageRuntime, threadIdRef, threadId }),
+  );
 };
 
 const MessageClientById = resource(useMessageClientById);
@@ -72,11 +79,24 @@ const useThreadClient = ({
     () => liveRef(() => runtime.getState()!.threadId),
     [runtime],
   );
+  const emitThreadEvent = (
+    event: "thread.cancelRun" | "thread.voiceStarted",
+  ) => {
+    emit(event, { threadId: runtime.getState()!.threadId });
+  };
+  const isSuggestion = useCallback(
+    (text: string) =>
+      runtime
+        .getState()!
+        .suggestions.some((suggestion) => suggestion.prompt === text),
+    [runtime],
+  );
 
   const composer = useClientResource(
     ComposerClient({
       runtime: runtime.composer,
       threadIdRef,
+      isSuggestion,
     }),
   );
   const suggestions = useClientResource(
@@ -84,11 +104,16 @@ const useThreadClient = ({
   );
   const messages = useClientLookup(
     runtimeState.messages.map((m) =>
-      withKey(m.id, MessageClientById({ runtime, id: m.id, threadIdRef }), [
-        runtime,
+      withKey(
         m.id,
-        threadIdRef,
-      ]),
+        MessageClientById({
+          runtime,
+          id: m.id,
+          threadIdRef,
+          threadId: runtimeState.threadId,
+        }),
+        [runtime, m.id, threadIdRef, runtimeState.threadId],
+      ),
     ),
   );
 
@@ -114,18 +139,41 @@ const useThreadClient = ({
     getState: () => state,
     composer: () => composer.methods,
     suggestions: () => suggestions.methods,
-    append: runtime.append,
+    append: (message) => {
+      const appended: Exclude<CreateAppendMessage, string> =
+        typeof message === "string"
+          ? { content: [{ type: "text", text: message }] }
+          : message;
+      if ((appended.role ?? "user") === "user") {
+        const text = appended.content
+          .map((part) => (part.type === "text" ? part.text : ""))
+          .join("");
+        emit("composer.send", {
+          threadId: runtime.getState()!.threadId,
+          chars: text.length,
+          attachments: appended.attachments?.length ?? 0,
+          ...(isSuggestion(text) ? { suggestion: true } : undefined),
+        });
+      }
+      runtime.append(message);
+    },
     deleteMessage: runtime.deleteMessage,
     startRun: runtime.startRun,
     resumeRun: runtime.resumeRun,
     importExternalState: runtime.importExternalState,
-    cancelRun: runtime.cancelRun,
+    cancelRun: () => {
+      if (runtimeState.isRunning) emitThreadEvent("thread.cancelRun");
+      runtime.cancelRun();
+    },
     getModelContext: runtime.getModelContext,
     export: runtime.export,
     import: runtime.import,
     reset: runtime.reset,
     stopSpeaking: runtime.stopSpeaking,
-    connectVoice: runtime.connectVoice,
+    connectVoice: () => {
+      runtime.connectVoice();
+      emitThreadEvent("thread.voiceStarted");
+    },
     disconnectVoice: runtime.disconnectVoice,
     getVoiceVolume: runtime.getVoiceVolume,
     subscribeVoiceVolume: runtime.subscribeVoiceVolume,

package/src/store/scopes/composer.ts

@@ -109,22 +109,34 @@ export type ComposerMeta = {
 
 export type ComposerEvents = {
   /**
-   * The user sent the composer contents. `messageId` is set when the send
-   * came from an edit composer.
+   * The user sent a message, from the composer or through `thread.append`.
+   * `messageId` is set when the send came from an edit composer.
    */
-  "composer.send": { threadId: string; messageId?: string };
+  "composer.send": {
+    threadId: string;
+    messageId?: string;
+    chars: number;
+    attachments: number;
+    suggestion?: boolean;
+  };
   /**
    * An attachment was added to the composer. `messageId` is set when the
    * attachment was added to an edit composer.
    */
-  "composer.attachmentAdd": { threadId: string; messageId?: string };
+  "composer.attachmentAdd": {
+    threadId: string;
+    messageId?: string;
+    contentType?: string;
+  };
   "composer.attachmentAddError": {
     threadId: string;
     messageId?: string;
     attachmentId?: string;
     reason: AttachmentAddErrorReason;
     message: string;
+    contentType?: string;
   };
+  "composer.cancel": { threadId: string };
 };
 
 export type ComposerClientSchema = {

package/src/store/scopes/message.ts

@@ -66,7 +66,21 @@ export type MessageMeta = {
   query: { type: "id"; id: string } | { type: "index"; index: number };
 };
 
+export type MessageEvents = {
+  /** The user asked for this assistant message to be generated again. */
+  "message.reload": { threadId: string; messageId: string };
+  /** The user moved to another branch at this message. */
+  "message.branchSwitched": { threadId: string; messageId: string };
+  /** The user copied this message. */
+  "message.copied": { threadId: string; messageId: string };
+  /** Speech synthesis started for this message. */
+  "message.speak": { threadId: string; messageId: string };
+  /** This message entered an error status. */
+  "message.error": { threadId: string; messageId: string; reason: "error" };
+};
+
 export type MessageClientSchema = {
   methods: MessageMethods;
   meta: MessageMeta;
+  events: MessageEvents;
 };

package/src/store/scopes/thread-list-item.ts

@@ -1,4 +1,7 @@
-import type { ThreadListItemRuntime } from "../../runtime/api/thread-list-item-runtime";
+import type {
+  ThreadListItemGenerateTitleOptions,
+  ThreadListItemRuntime,
+} from "../../runtime/api/thread-list-item-runtime";
 import type { ThreadListItemStatus } from "../../runtime/interfaces/thread-list-runtime-core";
 
 export type ThreadListItemState = {
@@ -26,7 +29,7 @@ export type ThreadListItemMethods = {
   archive(): void;
   unarchive(): void;
   delete(): void;
-  generateTitle(): void;
+  generateTitle(options?: ThreadListItemGenerateTitleOptions): void;
   initialize(): Promise<{ remoteId: string; externalId: string | undefined }>;
   detach(): void;
   __internal_getRuntime?(): ThreadListItemRuntime;

package/src/store/scopes/thread.ts

@@ -155,6 +155,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/tests/MessageRepository.test.ts

@@ -240,6 +240,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 +560,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 +1040,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,

package/src/tests/RemoteThreadListThreadListRuntimeCore-initialize-adapter-change.test.ts

@@ -0,0 +1,38 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+  createCore,
+  deferred,
+  makeAdapter,
+} from "./remote-thread-list-test-helpers";
+
+type InitializeResult = { remoteId: string; externalId: string };
+
+describe("RemoteThreadListThreadListRuntimeCore initialize", () => {
+  it("keeps the initialization task on the promoted slot when the adapter changes mid-flight", async () => {
+    const initializing = deferred<InitializeResult>();
+    const core = createCore(
+      makeAdapter({ initialize: vi.fn(() => initializing.promise) }),
+    );
+
+    await core.switchToNewThread();
+    const localId = core.newThreadId!;
+    const pending = core.initialize(localId);
+
+    // An adapter swap advances the generation without resetting the store, so
+    // the completion still applies its optimistic transform while `then`
+    // declines to reconcile against the retired adapter.
+    core.__internal_setOptions({
+      adapter: makeAdapter(),
+      runtimeHook: () => ({}) as never,
+    });
+
+    initializing.resolve({ remoteId: "remote-1", externalId: "external-1" });
+    await expect(pending).rejects.toThrow();
+
+    const item = core.getItemById(localId);
+    expect(item?.status).toBe("regular");
+    await expect(
+      item?.status === "new" ? undefined : item?.initializeTask,
+    ).resolves.toEqual({ remoteId: "remote-1", externalId: "external-1" });
+  });
+});