@assistant-ui/core 0.3.16 → 0.3.18

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/runtime-clients/useSubscribable.ts

@@ -2,7 +2,13 @@ import { useSyncExternalStore } from "react";
 import type { SubscribableWithState } from "../../subscribable/subscribable";
 
 export const useSubscribable = <T>(
-  subscribable: Omit<SubscribableWithState<T, any>, "path">,
+  subscribable: Omit<SubscribableWithState<T, any>, "path"> & {
+    getServerSnapshot?: () => T;
+  },
 ) => {
-  return useSyncExternalStore(subscribable.subscribe, subscribable.getState);
+  return useSyncExternalStore(
+    subscribable.subscribe,
+    subscribable.getState,
+    subscribable.getServerSnapshot,
+  );
 };

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/part.ts

@@ -29,10 +29,13 @@ export type PartMethods = {
   resumeToolCall(payload: unknown): void;
   /**
    * Respond to a server-side tool approval gate. The approval id is read from
-   * the part. Accepts a boolean decision or the id of one of the approval's
-   * options.
+   * the part. Accepts a boolean decision, the id of one of the approval's
+   * options, or a free-form answer when the request asks for one.
+   *
+   * Resolves once the runtime has accepted the response, and rejects when it
+   * could not be recorded, so a renderer can leave its controls retryable.
    */
-  respondToToolApproval(response: ToolApprovalResponse): void;
+  respondToToolApproval(response: ToolApprovalResponse): Promise<void>;
   __internal_getRuntime?(): MessagePartRuntime;
 };
 

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/subscribable/subscribable.test.ts

@@ -32,6 +32,30 @@ const createBinding = (initialState: TestState) => {
   };
 };
 
+// Mirrors runtime bindings such as `getThreadListState`, which assemble a new
+// object on every read while the underlying values keep their identity.
+const createRebuildingBinding = (initialState: TestState) => {
+  let state = initialState;
+  const subscribers = new Set<() => void>();
+
+  const binding: SubscribableWithState<TestState, null> = {
+    path: null,
+    getState: () => ({ ...state }),
+    subscribe: (callback) => {
+      subscribers.add(callback);
+      return () => subscribers.delete(callback);
+    },
+  };
+
+  return {
+    binding,
+    update(nextState: TestState) {
+      state = nextState;
+      for (const callback of subscribers) callback();
+    },
+  };
+};
+
 describe("ShallowMemoizeSubject", () => {
   it("notifies subscribers when a state key is removed", () => {
     const source = createBinding({
@@ -131,4 +155,30 @@ describe("LazyMemoizeSubject", () => {
     expect(subject.getState()).toEqual({ status: "ready" });
     reconnect();
   });
+
+  it("keeps one reference while unconnected and the state is unchanged", () => {
+    // A consumer that reads without subscribing (React's useSyncExternalStore
+    // calls getSnapshot outside the subscription) must not see a new object on
+    // every call: React compares snapshots with Object.is, so a fresh reference
+    // reads as a change and re-renders forever.
+    const source = createRebuildingBinding({ status: "ready" });
+    const subject = new LazyMemoizeSubject(source.binding);
+
+    const first = subject.getState();
+
+    expect(subject.getState()).toBe(first);
+    expect(subject.getState()).toBe(first);
+  });
+
+  it("returns a new reference once the unchanged state actually changes", () => {
+    const source = createRebuildingBinding({ status: "ready" });
+    const subject = new LazyMemoizeSubject(source.binding);
+
+    const first = subject.getState();
+    source.update({ status: "done" });
+
+    const second = subject.getState();
+    expect(second).not.toBe(first);
+    expect(second).toEqual({ status: "done" });
+  });
 });

package/src/subscribable/subscribable.ts

@@ -1,3 +1,4 @@
+import { shallowEqual } from "@assistant-ui/store/client";
 import type { Unsubscribe } from "../types/unsubscribe";
 import { notifyEventListeners } from "../utils/notify-event-listeners";
 
@@ -57,25 +58,10 @@ export const notifySubscribers = <TArgs extends unknown[]>(
   }
 };
 
-function shallowEqual<T extends object>(
-  objA: T | undefined,
-  objB: T | undefined,
-) {
-  if (objA === undefined && objB === undefined) return true;
-  if (objA === undefined) return false;
-  if (objB === undefined) return false;
-
-  const keysA = Object.keys(objA);
-  if (keysA.length !== Object.keys(objB).length) return false;
-
-  for (const key of keysA) {
-    const valueA = objA[key as keyof T];
-    const valueB = objB[key as keyof T];
-    if (!Object.is(valueA, valueB)) return false;
-  }
-
-  return true;
-}
+const shallowEqualOrUndefined = <T extends object>(
+  a: T | undefined,
+  b: T | undefined,
+) => (a === undefined || b === undefined ? a === b : shallowEqual(a, b));
 
 export class BaseSubscribable {
   private _subscribers = new Set<() => void>();
@@ -99,6 +85,32 @@ export class BaseSubscribable {
   }
 }
 
+export class WritableSubscribable<TState> extends BaseSubscribable {
+  private _state: TState;
+
+  constructor(state: TState) {
+    super();
+    this._state = state;
+    this.subscribe = this.subscribe.bind(this);
+    this.getState = this.getState.bind(this);
+    // Hydration has to agree with what the server rendered, so the server
+    // snapshot stays at the creation-time state rather than following writes.
+    this.getServerSnapshot = () => state;
+  }
+
+  public getState(): TState {
+    return this._state;
+  }
+
+  public getServerSnapshot: () => TState;
+
+  public setState(state: TState): void {
+    if (Object.is(state, this._state)) return;
+    this._state = state;
+    this._notifySubscribers();
+  }
+}
+
 // lazy connect/disconnect: only opens upstream subscription while it has subscribers
 export abstract class BaseSubject {
   private _subscriptions = new Set<(payload?: unknown) => void>();
@@ -168,7 +180,7 @@ export class ShallowMemoizeSubject<TState extends object, TPath>
   private _syncState() {
     const state = this.binding.getState();
     if (state === SKIP_UPDATE) return false;
-    if (shallowEqual(state, this._previousState)) return false;
+    if (shallowEqualOrUndefined(state, this._previousState)) return false;
     this._previousState = state;
     return true;
   }
@@ -209,7 +221,7 @@ export class LazyMemoizeSubject<TState extends object, TPath>
       if (
         newState !== SKIP_UPDATE &&
         (this._previousState === undefined ||
-          !shallowEqual(newState, this._previousState))
+          !shallowEqualOrUndefined(newState, this._previousState))
       ) {
         this._previousState = newState;
       }

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" });
+  });
+});