@assistant-ui/react 0.12.9 → 0.12.11

package/src/legacy-runtime/runtime-cores/assistant-transport/transport-scheduling.test.ts

@@ -0,0 +1,136 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { useRef } from "react";
+import { useCommandQueue } from "./commandQueue";
+import { useRunManager } from "./runManager";
+import type { AssistantTransportCommand } from "./types";
+
+const createMessageCommand = (id: string): AssistantTransportCommand => ({
+  type: "add-message",
+  message: {
+    role: "user",
+    parts: [{ type: "text", text: id }],
+  },
+  parentId: null,
+  sourceId: null,
+});
+
+const createDeferred = () => {
+  let resolve!: () => void;
+  const promise = new Promise<void>((res) => {
+    resolve = res;
+  });
+  return { promise, resolve };
+};
+
+const useTransportSchedulingHarness = (
+  opts: {
+    onRun?: (signal: AbortSignal) => Promise<void> | void;
+    onCancel?: (commands: AssistantTransportCommand[]) => void;
+  } = {},
+) => {
+  const commandQueueRef = useRef<ReturnType<typeof useCommandQueue> | null>(
+    null,
+  );
+  const runBatchesRef = useRef<AssistantTransportCommand[][]>([]);
+
+  const runManager = useRunManager({
+    onRun: async (signal) => {
+      const batch = commandQueueRef.current!.flush();
+      runBatchesRef.current.push(batch);
+      await opts.onRun?.(signal);
+    },
+    onCancel: () => {
+      const queue = commandQueueRef.current!;
+      const commands = [...queue.state.inTransit, ...queue.state.queued];
+      queue.reset();
+      opts.onCancel?.(commands);
+    },
+    onError: async () => {
+      // not needed for these contract tests
+    },
+  });
+
+  const commandQueue = useCommandQueue({
+    onQueue: () => runManager.schedule(),
+  });
+  commandQueueRef.current = commandQueue;
+
+  return {
+    commandQueue,
+    runManager,
+    runBatchesRef,
+  };
+};
+
+describe("assistant transport scheduling contracts", () => {
+  it("runs in single-flight mode and schedules exactly one follow-up run", async () => {
+    const gate = createDeferred();
+    const { result } = renderHook(() =>
+      useTransportSchedulingHarness({
+        onRun: () => gate.promise,
+      }),
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m1"));
+      result.current.commandQueue.enqueue(createMessageCommand("m2"));
+    });
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+    expect(result.current.runBatchesRef.current[0]).toHaveLength(2);
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("m3"));
+    });
+
+    await Promise.resolve();
+    expect(result.current.runBatchesRef.current).toHaveLength(1);
+
+    gate.resolve();
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(2);
+    });
+    expect(result.current.runBatchesRef.current[1]).toHaveLength(1);
+  });
+
+  it("cancel returns combined in-flight and queued commands", async () => {
+    const onCancel = vi.fn();
+    const { result } = renderHook(() =>
+      useTransportSchedulingHarness({
+        onRun: (signal) =>
+          new Promise<void>((_resolve, reject) => {
+            signal.addEventListener(
+              "abort",
+              () => reject(new Error("aborted")),
+              { once: true },
+            );
+          }),
+        onCancel,
+      }),
+    );
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("in-flight"));
+    });
+
+    await waitFor(() => {
+      expect(result.current.runBatchesRef.current).toHaveLength(1);
+    });
+
+    act(() => {
+      result.current.commandQueue.enqueue(createMessageCommand("queued"));
+      result.current.runManager.cancel();
+    });
+
+    await waitFor(() => {
+      expect(onCancel).toHaveBeenCalledTimes(1);
+    });
+    expect(onCancel.mock.calls[0]?.[0]).toHaveLength(2);
+  });
+});

package/src/legacy-runtime/runtime-cores/assistant-transport/types.ts

@@ -1,5 +1,5 @@
+import type { ThreadMessage } from "@assistant-ui/core";
 import { ReadonlyJSONValue } from "assistant-stream/utils";
-import { ThreadMessage } from "../../../types";
 import { AttachmentAdapter, ThreadHistoryAdapter } from "..";
 import { UserCommands } from "../../../augmentations";
 import type { ToolExecutionStatus } from "./useToolInvocations";
@@ -31,6 +31,8 @@ export type AssistantMessage = {
 export type AddMessageCommand = {
   readonly type: "add-message";
   readonly message: UserMessage | AssistantMessage;
+  readonly parentId: string | null;
+  readonly sourceId: string | null;
 };
 
 export type AddToolResultCommand = {
@@ -109,6 +111,9 @@ export type AssistantTransportOptions<T> = {
     updateState: (updater: (state: T) => T) => void;
     error?: Error;
   }) => void;
+  capabilities?: {
+    edit?: boolean;
+  };
   adapters?: {
     attachments?: AttachmentAdapter | undefined;
     history?: ThreadHistoryAdapter | undefined;

package/src/legacy-runtime/runtime-cores/assistant-transport/{useAssistantTransportRuntime.tsx → useAssistantTransportRuntime.ts}

@@ -1,11 +1,11 @@
 "use client";
 
+import type { AppendMessage } from "@assistant-ui/core";
 import {
   type ReadonlyJSONObject,
   type ReadonlyJSONValue,
   asAsyncIterableStream,
 } from "assistant-stream/utils";
-import { AppendMessage } from "../../../types";
 import { useExternalStoreRuntime } from "../external-store/useExternalStoreRuntime";
 import { AssistantRuntime } from "../../runtime/AssistantRuntime";
 import { AddToolResultOptions } from "../core";
@@ -29,12 +29,42 @@ import { useCommandQueue } from "./commandQueue";
 import { useRunManager } from "./runManager";
 import { useConvertedState } from "./useConvertedState";
 import { ToolExecutionStatus, useToolInvocations } from "./useToolInvocations";
-import { createRequestHeaders } from "./utils";
+import { createRequestHeaders } from "@assistant-ui/core";
 import { useRemoteThreadListRuntime } from "../remote-thread-list/useRemoteThreadListRuntime";
-import { InMemoryThreadListAdapter } from "../remote-thread-list/adapter/in-memory";
+import { InMemoryThreadListAdapter } from "@assistant-ui/core";
 import { useAui, useAuiState } from "@assistant-ui/store";
 import { UserExternalState } from "../../../augmentations";
 
+const convertAppendMessageToCommand = (
+  message: AppendMessage,
+): AddMessageCommand => {
+  if (message.role !== "user")
+    throw new Error("Only user messages are supported");
+
+  const parts: UserMessagePart[] = [];
+  const content = [
+    ...message.content,
+    ...(message.attachments?.flatMap((a) => a.content) ?? []),
+  ];
+  for (const contentPart of content) {
+    if (contentPart.type === "text") {
+      parts.push({ type: "text", text: contentPart.text });
+    } else if (contentPart.type === "image") {
+      parts.push({ type: "image", image: contentPart.image });
+    }
+  }
+
+  return {
+    type: "add-message",
+    message: {
+      role: "user",
+      parts,
+    },
+    parentId: message.parentId,
+    sourceId: message.sourceId,
+  };
+};
+
 const symbolAssistantTransportExtras = Symbol("assistant-transport-extras");
 type AssistantTransportExtras = {
   [symbolAssistantTransportExtras]: true;
@@ -74,12 +104,12 @@ export function useAssistantTransportState<T>(
 export function useAssistantTransportState<T>(
   selector: (state: UserExternalState) => T = (t) => t as T,
 ): T | UserExternalState {
-  return useAuiState(({ thread }) =>
-    selector(asAssistantTransportExtras(thread.extras).state),
+  return useAuiState((s) =>
+    selector(asAssistantTransportExtras(s.thread.extras).state),
   );
 }
 
-const useAssistantTransportThreadRuntime = <T,>(
+const useAssistantTransportThreadRuntime = <T>(
   options: AssistantTransportOptions<T>,
 ): AssistantRuntime => {
   const agentStateRef = useRef(options.initialState);
@@ -90,7 +120,7 @@ const useAssistantTransportThreadRuntime = <T,>(
     onQueue: () => runManager.schedule(),
   });
 
-  const threadId = useAuiState(({ threadListItem }) => threadListItem.remoteId);
+  const threadId = useAuiState((s) => s.threadListItem.remoteId);
 
   const runManager = useRunManager({
     onRun: async (signal: AbortSignal) => {
@@ -254,37 +284,17 @@ const useAssistantTransportThreadRuntime = <T,>(
       state: agentStateRef.current as UserExternalState,
     } satisfies AssistantTransportExtras,
     onNew: async (message: AppendMessage): Promise<void> => {
-      if (message.role !== "user")
-        throw new Error("Only user messages are supported");
-
-      // Store parentId for the request
       parentIdRef.current = message.parentId;
-
-      // Convert AppendMessage to AddMessageCommand
-      const parts: UserMessagePart[] = [];
-
-      const content = [
-        ...message.content,
-        ...(message.attachments?.flatMap((a) => a.content) ?? []),
-      ];
-      for (const contentPart of content) {
-        if (contentPart.type === "text") {
-          parts.push({ type: "text", text: contentPart.text });
-        } else if (contentPart.type === "image") {
-          parts.push({ type: "image", image: contentPart.image });
-        }
-      }
-
-      const command: AddMessageCommand = {
-        type: "add-message",
-        message: {
-          role: "user",
-          parts,
-        },
-      };
-
+      const command = convertAppendMessageToCommand(message);
       commandQueue.enqueue(command);
     },
+    ...(options.capabilities?.edit && {
+      onEdit: async (message: AppendMessage): Promise<void> => {
+        parentIdRef.current = message.parentId;
+        const command = convertAppendMessageToCommand(message);
+        commandQueue.enqueue(command);
+      },
+    }),
     onCancel: async () => {
       runManager.cancel();
       await toolInvocations.abort();
@@ -330,7 +340,7 @@ const useAssistantTransportThreadRuntime = <T,>(
 /**
  * @alpha This is an experimental API that is subject to change.
  */
-export const useAssistantTransportRuntime = <T,>(
+export const useAssistantTransportRuntime = <T>(
   options: AssistantTransportOptions<T>,
 ): AssistantRuntime => {
   const runtime = useRemoteThreadListRuntime({

package/src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.test.ts

@@ -0,0 +1,223 @@
+// @vitest-environment jsdom
+
+import type { ThreadAssistantMessage } from "@assistant-ui/core";
+import type { Tool } from "assistant-stream";
+import { act, renderHook, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import type { AssistantTransportState, ToolExecutionStatus } from "./types";
+import { useToolInvocations } from "./useToolInvocations";
+
+const createState = (
+  messages: ThreadAssistantMessage[],
+): AssistantTransportState => ({
+  messages,
+  isRunning: true,
+});
+
+const createAssistantMessage = (
+  argsText: string,
+  args: Record<string, unknown>,
+): ThreadAssistantMessage => ({
+  id: "m-1",
+  role: "assistant",
+  createdAt: new Date(),
+  status: { type: "requires-action", reason: "tool-calls" },
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+  content: [
+    {
+      type: "tool-call",
+      toolCallId: "tool-1",
+      toolName: "weatherSearch",
+      args,
+      argsText,
+    },
+  ],
+});
+
+describe("useToolInvocations", () => {
+  it("does not crash when tool argsText rewrites a previously streamed value", async () => {
+    const execute = vi.fn(async () => ({ forecast: "ok" }));
+    const getTools = () => ({
+      weatherSearch: {
+        execute,
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+    const setToolStatuses = vi.fn();
+
+    const { rerender } = renderHook(
+      ({ state }: { state: AssistantTransportState }) =>
+        useToolInvocations({
+          state,
+          getTools,
+          onResult,
+          setToolStatuses,
+        }),
+      {
+        initialProps: {
+          state: createState([]),
+        },
+      },
+    );
+
+    expect(() => {
+      act(() => {
+        rerender({
+          state: createState([
+            createAssistantMessage('{"query":"London","longitude":0', {
+              query: "London",
+              longitude: 0,
+            }),
+          ]),
+        });
+      });
+    }).not.toThrow();
+
+    expect(() => {
+      act(() => {
+        rerender({
+          state: createState([
+            createAssistantMessage('{"query":"London","longitude":-0.125', {
+              query: "London",
+              longitude: -0.125,
+            }),
+          ]),
+        });
+      });
+    }).not.toThrow();
+
+    act(() => {
+      rerender({
+        state: createState([
+          createAssistantMessage(
+            '{"query":"London","longitude":-0.125,"latitude":51.5072}',
+            { query: "London", longitude: -0.125, latitude: 51.5072 },
+          ),
+        ]),
+      });
+    });
+
+    await waitFor(() => {
+      expect(execute).toHaveBeenCalledTimes(1);
+    });
+
+    expect(execute).toHaveBeenCalledWith(
+      {
+        query: "London",
+        longitude: -0.125,
+        latitude: 51.5072,
+      },
+      expect.objectContaining({ toolCallId: "tool-1" }),
+    );
+
+    await waitFor(() => {
+      expect(onResult).toHaveBeenCalledTimes(1);
+    });
+    expect(onResult).toHaveBeenCalledWith(
+      expect.objectContaining({
+        type: "add-tool-result",
+        toolCallId: "tool-1",
+        toolName: "weatherSearch",
+      }),
+    );
+  });
+
+  it("keeps logical toolCallId mapping through reset while abort settles", async () => {
+    const execute = vi.fn(
+      async () =>
+        await new Promise(() => {
+          // never resolves: reset() should cancel this call
+        }),
+    );
+    const getTools = () => ({
+      weatherSearch: {
+        execute,
+      } satisfies Tool,
+    });
+    const onResult = vi.fn();
+
+    let statuses: Record<string, ToolExecutionStatus> = {};
+    const setToolStatuses = vi.fn(
+      (
+        updater:
+          | Record<string, ToolExecutionStatus>
+          | ((
+              prev: Record<string, ToolExecutionStatus>,
+            ) => Record<string, ToolExecutionStatus>),
+      ) => {
+        statuses =
+          typeof updater === "function" ? updater(statuses) : { ...updater };
+      },
+    );
+
+    const { result, rerender } = renderHook(
+      ({ state }: { state: AssistantTransportState }) =>
+        useToolInvocations({
+          state,
+          getTools,
+          onResult,
+          setToolStatuses,
+        }),
+      {
+        initialProps: {
+          state: createState([]),
+        },
+      },
+    );
+
+    act(() => {
+      rerender({
+        state: createState([
+          createAssistantMessage('{"query":"London","longitude":0', {
+            query: "London",
+            longitude: 0,
+          }),
+        ]),
+      });
+    });
+
+    act(() => {
+      rerender({
+        state: createState([
+          createAssistantMessage('{"query":"London","longitude":-0.125', {
+            query: "London",
+            longitude: -0.125,
+          }),
+        ]),
+      });
+    });
+
+    act(() => {
+      rerender({
+        state: createState([
+          createAssistantMessage(
+            '{"query":"London","longitude":-0.125,"latitude":51.5072}',
+            { query: "London", longitude: -0.125, latitude: 51.5072 },
+          ),
+        ]),
+      });
+    });
+
+    await waitFor(() => {
+      expect(execute).toHaveBeenCalledTimes(1);
+    });
+    await waitFor(() => {
+      expect(statuses["tool-1"]).toEqual({ type: "executing" });
+    });
+
+    act(() => {
+      result.current.reset();
+    });
+
+    await waitFor(() => {
+      expect(statuses).toEqual({});
+    });
+    expect(Object.keys(statuses)).not.toContain("tool-1:rewrite:0");
+  });
+});