@assistant-ui/core 0.3.2 → 0.3.4

package/src/react/primitive-hooks/useSuggestionTrigger.test.ts

@@ -0,0 +1,148 @@
+/** @vitest-environment jsdom */
+import { cleanup, renderHook } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => {
+  const append = vi.fn();
+  const setText = vi.fn();
+
+  const state = {
+    thread: {
+      isDisabled: false,
+      isRunning: false,
+      capabilities: { queue: false },
+    },
+  };
+  const composerState = {
+    text: "",
+    runConfig: { custom: { model: "gpt-test" } },
+  };
+
+  return {
+    append,
+    setText,
+    state,
+    composerState,
+    aui: {
+      thread: {
+        getState: () => state.thread,
+        append,
+      },
+      composer: {
+        getState: () => composerState,
+        setText,
+      },
+    },
+  };
+});
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAui: () => mocks.aui,
+  useAuiState: ((selector: (state: typeof mocks.state) => unknown) =>
+    selector(mocks.state)) as typeof import("@assistant-ui/store").useAuiState,
+}));
+
+import { useSuggestionTrigger } from "./useSuggestionTrigger";
+
+afterEach(() => {
+  cleanup();
+  vi.clearAllMocks();
+  mocks.state.thread.isDisabled = false;
+  mocks.state.thread.isRunning = false;
+  mocks.state.thread.capabilities = { queue: false };
+  mocks.composerState.text = "";
+});
+
+describe("useSuggestionTrigger", () => {
+  it("appends the prompt and clears the composer when sending while idle", () => {
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", send: true }),
+    );
+
+    result.current.trigger();
+
+    expect(result.current.disabled).toBe(false);
+    expect(mocks.append).toHaveBeenCalledWith({
+      content: [{ type: "text", text: "Hello" }],
+      runConfig: { custom: { model: "gpt-test" } },
+    });
+    expect(mocks.setText).toHaveBeenCalledWith("");
+  });
+
+  it("disables and no-ops when sending while running without queue support", () => {
+    mocks.state.thread.isRunning = true;
+    mocks.composerState.text = "my draft";
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", send: true }),
+    );
+
+    result.current.trigger();
+
+    expect(result.current.disabled).toBe(true);
+    expect(mocks.append).not.toHaveBeenCalled();
+    expect(mocks.setText).not.toHaveBeenCalled();
+  });
+
+  it("queues without touching the composer while running when the thread supports queueing", () => {
+    mocks.state.thread.isRunning = true;
+    mocks.state.thread.capabilities = { queue: true };
+    mocks.composerState.text = "my draft";
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", send: true }),
+    );
+
+    result.current.trigger();
+
+    expect(result.current.disabled).toBe(false);
+    expect(mocks.append).toHaveBeenCalledWith({
+      content: [{ type: "text", text: "Hello" }],
+      runConfig: { custom: { model: "gpt-test" } },
+    });
+    expect(mocks.setText).not.toHaveBeenCalled();
+  });
+
+  it("replaces the composer text when send is false, even while running", () => {
+    mocks.state.thread.isRunning = true;
+    mocks.composerState.text = "my draft";
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello" }),
+    );
+
+    result.current.trigger();
+
+    expect(result.current.disabled).toBe(false);
+    expect(mocks.append).not.toHaveBeenCalled();
+    expect(mocks.setText).toHaveBeenCalledWith("Hello");
+  });
+
+  it("appends to the composer text when clearComposer is false", () => {
+    mocks.composerState.text = "my draft";
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", clearComposer: false }),
+    );
+
+    result.current.trigger();
+
+    expect(mocks.setText).toHaveBeenCalledWith("my draft Hello");
+  });
+
+  it("inserts the prompt alone when the composer is empty and clearComposer is false", () => {
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello", clearComposer: false }),
+    );
+
+    result.current.trigger();
+
+    expect(mocks.setText).toHaveBeenCalledWith("Hello");
+  });
+
+  it("disables when the thread is disabled", () => {
+    mocks.state.thread.isDisabled = true;
+    const { result } = renderHook(() =>
+      useSuggestionTrigger({ prompt: "Hello" }),
+    );
+
+    expect(result.current.disabled).toBe(true);
+  });
+});

package/src/react/primitive-hooks/useSuggestionTrigger.ts

@@ -13,18 +13,24 @@ export const useSuggestionTrigger = ({
   clearComposer = true,
 }: UseSuggestionTriggerOptions) => {
   const aui = useAui();
-  const disabled = useAuiState((s) => s.thread.isDisabled);
   const resolvedSend = send ?? false;
+  const disabled = useAuiState(
+    (s) =>
+      s.thread.isDisabled ||
+      (resolvedSend && s.thread.isRunning && !s.thread.capabilities.queue),
+  );
 
   const trigger = useCallback(() => {
-    const isRunning = aui.thread.getState().isRunning;
+    if (resolvedSend) {
+      const { isRunning, capabilities } = aui.thread.getState();
+      if (isRunning && !capabilities.queue) return;
 
-    if (resolvedSend && !isRunning) {
       aui.thread.append({
         content: [{ type: "text", text: prompt }],
         runConfig: aui.composer.getState().runConfig,
       });
-      if (clearComposer) {
+      // A queued send must not clear the draft the user is still composing.
+      if (clearComposer && !isRunning) {
         aui.composer.setText("");
       }
     } else {

package/src/react/primitives/message/MessageGroupedParts.tsx

@@ -8,6 +8,7 @@ import type {
   MessagePartStatus,
   ToolCallMessagePartStatus,
 } from "../../../types/message";
+import { getGroupStatus } from "../../../utils/getGroupStatus";
 import {
   buildGroupTree,
   GROUPBY_MEMO_KEY,
@@ -21,7 +22,8 @@ export namespace MessagePrimitiveGroupedParts {
    * A coalesced group of adjacent parts. Surfaced through the same
    * `{ part }` channel as a leaf {@link EnrichedPartState} so consumers
    * dispatch on a single `switch (part.type)`. `type` is the group key
-   * (always `"group-…"`); `status` mirrors the last contained part.
+   * (always `"group-…"`); `status` is running when any contained part runs,
+   * otherwise it mirrors the last contained part.
    */
   export type GroupPart<TKey extends `group-${string}` = `group-${string}`> = {
     readonly type: TKey;
@@ -133,8 +135,6 @@ export namespace MessagePrimitiveGroupedParts {
   };
 }
 
-const COMPLETE_STATUS: MessagePartStatus = Object.freeze({ type: "complete" });
-
 const shouldShowIndicator = (
   mode: MessagePrimitiveGroupedParts.IndicatorMode,
   parts: readonly PartState[],
@@ -194,7 +194,7 @@ const renderNode = <TKey extends `group-${string}`>(
     );
   }
 
-  const status = parts[node.indices.at(-1)!]?.status ?? COMPLETE_STATUS;
+  const status = getGroupStatus(parts, node.indices);
   const groupPart: MessagePrimitiveGroupedParts.GroupPart<TKey> = {
     type: node.key as TKey,
     status,

package/src/react/primitives/message/MessageParts.tsx

@@ -214,7 +214,12 @@ export namespace MessagePrimitiveParts {
     Image?: ImageMessagePartComponent | undefined;
     /** Component for rendering file content */
     File?: FileMessagePartComponent | undefined;
-    /** Component for rendering audio content (experimental) */
+    /**
+     * Component for rendering audio content.
+     *
+     * @deprecated Render audio through the `File` slot instead, branching on an
+     * `audio/*` mime type.
+     */
     Unstable_Audio?: Unstable_AudioMessagePartComponent | undefined;
     /** Configuration for data part rendering */
     data?: DataConfig | undefined;

package/src/react/providers/ReadonlyThreadProvider.tsx

@@ -26,6 +26,7 @@ const READONLY_THREAD_LIST_ITEM: ThreadListItemState = Object.freeze({
   remoteId: undefined,
   externalId: undefined,
   isMain: true,
+  isRunning: false,
   status: "regular" as const,
   title: undefined,
 });

package/src/react/runtimes/RemoteThreadListHookInstanceManager.running.test.tsx

@@ -0,0 +1,157 @@
+import { describe, expect, it, vi } from "vitest";
+import type { ThreadListRuntimeCore } from "../../runtime/interfaces/thread-list-runtime-core";
+import type { ThreadRuntimeCore } from "../../runtime/interfaces/thread-runtime-core";
+import type { ThreadMessage } from "../../types/message";
+import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
+
+const makeRuntime = (
+  initial: Partial<Pick<ThreadRuntimeCore, "isRunning" | "messages">> = {},
+) => {
+  const subscribers = new Set<() => void>();
+  const runtime = {
+    isRunning: initial.isRunning,
+    messages: initial.messages ?? [],
+    subscribe: (callback: () => void) => {
+      subscribers.add(callback);
+      return () => subscribers.delete(callback);
+    },
+  } as unknown as ThreadRuntimeCore & {
+    isRunning: boolean | undefined;
+    messages: readonly ThreadMessage[];
+  };
+
+  return {
+    runtime,
+    subscriberCount: () => subscribers.size,
+    setRunning: (isRunning: boolean | undefined) => {
+      runtime.isRunning = isRunning;
+      for (const callback of subscribers) callback();
+    },
+  };
+};
+
+const makeManager = () =>
+  new RemoteThreadListHookInstanceManager(
+    () => ({}) as never,
+    {} as ThreadListRuntimeCore,
+  );
+
+// mirrors what the React binder does on every publication
+const publish = (
+  manager: RemoteThreadListHookInstanceManager,
+  threadId: string,
+  runtime: ThreadRuntimeCore,
+) => {
+  const internals = manager as unknown as {
+    instances: Map<string, { generation: number }>;
+    _publishThreadRuntime: (
+      threadId: string,
+      runtime: ThreadRuntimeCore,
+      generation: number,
+    ) => void;
+  };
+  const generation = internals.instances.get(threadId)!.generation;
+  internals._publishThreadRuntime(threadId, runtime, generation);
+};
+
+const start = (manager: RemoteThreadListHookInstanceManager, id: string) => {
+  manager.startThreadRuntime(id).catch(() => {});
+};
+
+describe("RemoteThreadListHookInstanceManager run tracking", () => {
+  it("reports a thread with no attached runtime as not running", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(false);
+    expect(manager.__internal_isThreadRunning("never-started")).toBe(false);
+  });
+
+  it("adopts the run state of the published runtime", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const { runtime } = makeRuntime({ isRunning: true });
+
+    publish(manager, "thread-1", runtime);
+
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(true);
+  });
+
+  it("falls back to the trailing assistant message when the runtime does not track runs", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const { runtime } = makeRuntime({
+      messages: [
+        { role: "assistant", status: { type: "running" } },
+      ] as unknown as readonly ThreadMessage[],
+    });
+
+    publish(manager, "thread-1", runtime);
+
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(true);
+  });
+
+  it("tracks a thread the user has switched away from", () => {
+    const manager = makeManager();
+    start(manager, "background");
+    start(manager, "main");
+    const background = makeRuntime({ isRunning: false });
+    publish(manager, "background", background.runtime);
+    publish(manager, "main", makeRuntime({ isRunning: false }).runtime);
+
+    background.setRunning(true);
+
+    expect(manager.__internal_isThreadRunning("background")).toBe(true);
+    expect(manager.__internal_isThreadRunning("main")).toBe(false);
+  });
+
+  it("notifies the thread list only when the thread crosses the running boundary", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const thread = makeRuntime({ isRunning: false });
+    publish(manager, "thread-1", thread.runtime);
+
+    const onChange = vi.fn();
+    manager.__internal_subscribeRunningChanged(onChange);
+
+    thread.setRunning(false);
+    expect(onChange).not.toHaveBeenCalled();
+
+    thread.setRunning(true);
+    expect(onChange).toHaveBeenCalledTimes(1);
+
+    thread.setRunning(true);
+    expect(onChange).toHaveBeenCalledTimes(1);
+
+    thread.setRunning(false);
+    expect(onChange).toHaveBeenCalledTimes(2);
+  });
+
+  it("moves tracking to the runtime a restart publishes", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const before = makeRuntime({ isRunning: true });
+    publish(manager, "thread-1", before.runtime);
+
+    const after = makeRuntime({ isRunning: false });
+    publish(manager, "thread-1", after.runtime);
+
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(false);
+    expect(before.subscriberCount()).toBe(0);
+
+    after.setRunning(true);
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(true);
+  });
+
+  it("releases the run subscription when the thread runtime stops", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const thread = makeRuntime({ isRunning: true });
+    publish(manager, "thread-1", thread.runtime);
+
+    manager.stopThreadRuntime("thread-1");
+
+    expect(thread.subscriberCount()).toBe(0);
+    expect(manager.__internal_isThreadRunning("thread-1")).toBe(false);
+  });
+});

package/src/react/runtimes/RemoteThreadListHookInstanceManager.test.tsx

@@ -16,3 +16,163 @@ describe("RemoteThreadListHookInstanceManager", () => {
     );
   });
 });
+
+describe("RemoteThreadListHookInstanceManager.__internal_restartThreadRuntime", () => {
+  const makeManager = () =>
+    new RemoteThreadListHookInstanceManager(
+      () => ({}) as never,
+      {} as ThreadListRuntimeCore,
+    );
+
+  // no React binder attaches a runtime in these tests, so the returned promises
+  // stay pending or reject on stop; neither is what is under test here
+  const start = (manager: RemoteThreadListHookInstanceManager, id: string) => {
+    manager.startThreadRuntime(id).catch(() => {});
+  };
+  const restart = (
+    manager: RemoteThreadListHookInstanceManager,
+    id: string,
+  ) => {
+    manager.__internal_restartThreadRuntime(id).catch(() => {});
+  };
+
+  type InstanceInternals = {
+    instances: Map<
+      string,
+      {
+        runtime?: unknown;
+        publishedGeneration?: number;
+        generation: number;
+      }
+    >;
+    _notifySubscribers: () => void;
+  };
+  const internalsOf = (manager: RemoteThreadListHookInstanceManager) =>
+    manager as unknown as InstanceInternals;
+
+  const renderedKeys = (manager: RemoteThreadListHookInstanceManager) =>
+    Array.from(internalsOf(manager).instances.entries()).map(
+      ([id, { generation }]) => `${id}:${generation}`,
+    );
+
+  it("changes the binder key so React remounts the runtime hook", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const before = renderedKeys(manager);
+
+    restart(manager, "thread-1");
+
+    expect(renderedKeys(manager)).not.toEqual(before);
+    expect(renderedKeys(manager)).toEqual(["thread-1:1"]);
+  });
+
+  it("keeps the thread rendered across the restart, unlike stop", () => {
+    const restarted = makeManager();
+    start(restarted, "thread-1");
+    restart(restarted, "thread-1");
+
+    const stopped = makeManager();
+    start(stopped, "thread-1");
+    stopped.stopThreadRuntime("thread-1");
+
+    expect(renderedKeys(restarted)).toHaveLength(1);
+    expect(renderedKeys(stopped)).toHaveLength(0);
+  });
+
+  it("starts the runtime when the thread is not alive yet", () => {
+    const manager = makeManager();
+
+    restart(manager, "thread-1");
+
+    expect(renderedKeys(manager)).toEqual(["thread-1:0"]);
+  });
+
+  it("stop then start in one tick yields a fresh key, so the old still-mounted binder cannot satisfy the new start", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const before = renderedKeys(manager);
+
+    manager.stopThreadRuntime("thread-1");
+    start(manager, "thread-1");
+
+    expect(renderedKeys(manager)).not.toEqual(before);
+    expect(renderedKeys(manager)).toEqual(["thread-1:1"]);
+  });
+
+  // simulates the binder's updateRuntime: publish a runtime for the
+  // instance's current generation (or an explicit stale one)
+  const publish = (
+    manager: RemoteThreadListHookInstanceManager,
+    id: string,
+    runtime: unknown,
+    options?: { generation?: number },
+  ) => {
+    const instance = internalsOf(manager).instances.get(id)!;
+    instance.runtime = runtime;
+    instance.publishedGeneration = options?.generation ?? instance.generation;
+    internalsOf(manager)._notifySubscribers();
+  };
+
+  it("does not settle with the pre-restart runtime; only the incoming binder's publication resolves it", async () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    publish(manager, "thread-1", { tag: "pre-reload-runtime" });
+
+    let settledWith = "NOT_SETTLED";
+    manager.__internal_restartThreadRuntime("thread-1").then((r) => {
+      settledWith = (r as { tag: string }).tag;
+    });
+
+    // the outgoing runtime rides across the restart (stays readable)…
+    expect(manager.getThreadRuntimeCore("thread-1")).toEqual({
+      tag: "pre-reload-runtime",
+    });
+    // …but must not count as attached
+    await Promise.resolve();
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(settledWith).toBe("NOT_SETTLED");
+
+    publish(manager, "thread-1", { tag: "post-reload-runtime" });
+    await Promise.resolve();
+    expect(settledWith).toBe("post-reload-runtime");
+  });
+
+  it("keeps the outgoing runtime readable while the restart promise is pending", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    publish(manager, "thread-1", { tag: "pre-reload-runtime" });
+
+    restart(manager, "thread-1");
+
+    expect(manager.getThreadRuntimeCore("thread-1")).toEqual({
+      tag: "pre-reload-runtime",
+    });
+  });
+
+  it("a stale-generation publication does not resolve the restart promise", async () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    publish(manager, "thread-1", { tag: "pre-reload-runtime" });
+
+    let settled = false;
+    manager.__internal_restartThreadRuntime("thread-1").then(() => {
+      settled = true;
+    });
+
+    // an outgoing binder re-publishing for its old generation (e.g. a late
+    // outerSubscribe callback) must not count as the new attachment
+    const staleGeneration =
+      internalsOf(manager).instances.get("thread-1")!.generation - 1;
+    publish(
+      manager,
+      "thread-1",
+      { tag: "stale-publication" },
+      { generation: staleGeneration },
+    );
+
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(settled).toBe(false);
+  });
+});