@assistant-ui/react 0.14.9 → 0.14.11

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

@@ -1,892 +0,0 @@
-// @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 } from "./types";
-import {
-  type ToolExecutionStatus,
-  useToolInvocations,
-} from "./useToolInvocations";
-import type {
-  ReadonlyJSONObject,
-  ReadonlyJSONValue,
-} from "assistant-stream/utils";
-
-const createState = (
-  messages: ThreadAssistantMessage[],
-  isRunning: boolean = true,
-): AssistantTransportState => ({
-  messages,
-  isRunning,
-});
-
-const createAssistantMessage = (
-  argsText: string,
-  args: Record<string, unknown>,
-  options?: {
-    result?: ReadonlyJSONValue;
-    isError?: boolean;
-    toolCallId?: string;
-    toolName?: string;
-    nestedMessages?: ThreadAssistantMessage[];
-  },
-): 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: options?.toolCallId ?? "tool-1",
-      toolName: options?.toolName ?? "weatherSearch",
-      args: args as ReadonlyJSONObject,
-      argsText,
-      ...(options?.result !== undefined && { result: options.result }),
-      ...(options?.isError !== undefined && { isError: options.isError }),
-      ...(options?.nestedMessages && { messages: options.nestedMessages }),
-    },
-  ],
-});
-
-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: {
-        parameters: { type: "object", properties: {} },
-        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: {
-        parameters: { type: "object", properties: {} },
-        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");
-  });
-
-  it("does not execute tool calls loaded asynchronously with existing results", async () => {
-    const execute = vi.fn(async () => ({ forecast: "ok" }));
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        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([]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"London"}',
-            { query: "London" },
-            { result: { source: "history" } },
-          ),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(execute).not.toHaveBeenCalled();
-      expect(onResult).not.toHaveBeenCalled();
-    });
-  });
-
-  it("does not re-execute asynchronously loaded resolved tool calls after reset", async () => {
-    const execute = vi.fn(async () => ({ forecast: "ok" }));
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        execute,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    const { result, rerender } = renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage('{"query":"London"}', { query: "London" }),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(execute).toHaveBeenCalledTimes(1);
-    });
-
-    act(() => {
-      result.current.reset();
-    });
-
-    await act(async () => {
-      await Promise.resolve();
-    });
-
-    act(() => {
-      rerender({
-        state: createState([]),
-      });
-    });
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"London"}',
-            { query: "London" },
-            { result: { source: "history" } },
-          ),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(execute).toHaveBeenCalledTimes(1);
-      expect(onResult).toHaveBeenCalledTimes(1);
-    });
-  });
-
-  it("still processes nested unresolved tool calls when the parent tool call is already resolved", async () => {
-    const executeParent = vi.fn(async () => ({ scope: "parent" }));
-    const executeChild = vi.fn(async () => ({ scope: "child" }));
-    const getTools = () => ({
-      resolvedOnly: {
-        parameters: { type: "object", properties: {} },
-        execute: executeParent,
-      } satisfies Tool,
-      childTool: {
-        parameters: { type: "object", properties: {} },
-        execute: executeChild,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    const nestedMessage = createAssistantMessage(
-      '{"query":"nested"}',
-      { query: "nested" },
-      {
-        toolCallId: "tool-child",
-        toolName: "childTool",
-      },
-    );
-
-    const { rerender } = renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"parent"}',
-            { query: "parent" },
-            {
-              result: { source: "history" },
-              toolName: "resolvedOnly",
-              nestedMessages: [nestedMessage],
-            },
-          ),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(executeParent).not.toHaveBeenCalled();
-      expect(executeChild).toHaveBeenCalledTimes(1);
-    });
-  });
-
-  it("does not close args stream early for non-executable tool snapshots", () => {
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
-
-    try {
-      const { rerender } = renderHook(
-        ({ state }: { state: AssistantTransportState }) =>
-          useToolInvocations({
-            state,
-            getTools,
-            onResult,
-            setToolStatuses,
-          }),
-        {
-          initialProps: {
-            state: createState([]),
-          },
-        },
-      );
-
-      act(() => {
-        rerender({
-          state: createState([createAssistantMessage("{}", {})]),
-        });
-      });
-
-      act(() => {
-        rerender({
-          state: createState([
-            createAssistantMessage('{"title":"Weekly"', {
-              title: "Weekly",
-            }),
-          ]),
-        });
-      });
-
-      act(() => {
-        rerender({
-          state: createState([
-            createAssistantMessage('{"title":"Weekly","columns":["name"]}', {
-              title: "Weekly",
-              columns: ["name"],
-            }),
-          ]),
-        });
-      });
-
-      expect(warnSpy).not.toHaveBeenCalledWith(
-        "argsText updated after controller was closed:",
-        expect.anything(),
-      );
-      expect(warnSpy).not.toHaveBeenCalledWith(
-        "argsText updated after controller was closed, restarting tool args stream:",
-        expect.anything(),
-      );
-      expect(onResult).not.toHaveBeenCalled();
-    } finally {
-      warnSpy.mockRestore();
-    }
-  });
-
-  it("closes non-executable complete args stream after run settles", () => {
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-    const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
-
-    try {
-      const { rerender } = renderHook(
-        ({ state }: { state: AssistantTransportState }) =>
-          useToolInvocations({
-            state,
-            getTools,
-            onResult,
-            setToolStatuses,
-          }),
-        {
-          initialProps: {
-            state: createState([]),
-          },
-        },
-      );
-
-      act(() => {
-        rerender({
-          state: createState(
-            [
-              createAssistantMessage('{"title":"Weekly"}', {
-                title: "Weekly",
-              }),
-            ],
-            true,
-          ),
-        });
-      });
-
-      act(() => {
-        rerender({
-          state: createState(
-            [
-              createAssistantMessage('{"title":"Weekly"}', {
-                title: "Weekly",
-              }),
-            ],
-            false,
-          ),
-        });
-      });
-
-      act(() => {
-        rerender({
-          state: createState(
-            [
-              createAssistantMessage('{"title":"Weekly","columns":["name"]}', {
-                title: "Weekly",
-                columns: ["name"],
-              }),
-            ],
-            false,
-          ),
-        });
-      });
-
-      expect(warnSpy).toHaveBeenCalledWith(
-        "argsText updated after controller was closed, restarting tool args stream:",
-        expect.objectContaining({
-          previous: '{"title":"Weekly"}',
-          next: '{"title":"Weekly","columns":["name"]}',
-        }),
-      );
-      expect(onResult).not.toHaveBeenCalled();
-    } finally {
-      warnSpy.mockRestore();
-    }
-  });
-
-  it("handles backend result when equivalent complete argsText reorders keys", async () => {
-    let resolveExecute: ((value: unknown) => void) | undefined;
-    const execute = vi.fn(
-      () =>
-        new Promise<unknown>((resolve) => {
-          resolveExecute = resolve;
-        }),
-    );
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        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([]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage('{"a":1,"b":2}', {
-            a: 1,
-            b: 2,
-          }),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(execute).toHaveBeenCalledTimes(1);
-    });
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"b":2,"a":1}',
-            {
-              a: 1,
-              b: 2,
-            },
-            {
-              result: { source: "backend" },
-            },
-          ),
-        ]),
-      });
-    });
-
-    await act(async () => {
-      resolveExecute?.({ source: "client" });
-    });
-
-    await waitFor(() => {
-      expect(onResult).not.toHaveBeenCalled();
-    });
-  });
-
-  it("fires streamCall for already-resolved tool calls loaded after the initial snapshot", async () => {
-    const execute = vi.fn(async () => ({ forecast: "ok" }));
-    const streamCall = vi.fn();
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        execute,
-        streamCall,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    const { rerender } = renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"London"}',
-            { query: "London" },
-            { result: { source: "history" } },
-          ),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(streamCall).toHaveBeenCalledTimes(1);
-    });
-
-    const [reader] = streamCall.mock.calls[0]!;
-    await expect(reader.args.get("query")).resolves.toBe("London");
-    const response = await reader.response.get();
-    expect(response.result).toEqual({ source: "history" });
-
-    expect(execute).not.toHaveBeenCalled();
-    expect(onResult).not.toHaveBeenCalled();
-    expect(setToolStatuses).not.toHaveBeenCalled();
-  });
-
-  it("does not fire streamCall for tool calls present in the initial snapshot", async () => {
-    const streamCall = vi.fn();
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        streamCall,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([
-            createAssistantMessage(
-              '{"query":"London"}',
-              { query: "London" },
-              { result: { source: "history" } },
-            ),
-          ]),
-        },
-      },
-    );
-
-    await act(async () => {});
-    expect(streamCall).not.toHaveBeenCalled();
-    expect(onResult).not.toHaveBeenCalled();
-  });
-
-  it("promotes an in-progress tool call from the initial snapshot when it changes", async () => {
-    const execute = vi.fn(async () => ({ forecast: "ok" }));
-    const streamCall = vi.fn();
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        execute,
-        streamCall,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    const { rerender } = renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([
-            createAssistantMessage('{"query":"Lon', { query: "Lon" }),
-          ]),
-        },
-      },
-    );
-
-    await act(async () => {});
-    expect(streamCall).not.toHaveBeenCalled();
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"London"}',
-            { query: "London" },
-            { result: { source: "history" } },
-          ),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(streamCall).toHaveBeenCalledTimes(1);
-    });
-
-    const [reader] = streamCall.mock.calls[0]!;
-    await expect(reader.args.get("query")).resolves.toBe("London");
-    const response = await reader.response.get();
-    expect(response.result).toEqual({ source: "history" });
-
-    expect(execute).not.toHaveBeenCalled();
-    expect(onResult).not.toHaveBeenCalled();
-  });
-
-  it("does not re-fire streamCall when an initial-snapshot tool call is unchanged in later snapshots", async () => {
-    const streamCall = vi.fn();
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        streamCall,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    const { rerender } = renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([
-            createAssistantMessage(
-              '{"query":"London"}',
-              { query: "London" },
-              { result: { source: "history" } },
-            ),
-          ]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"London"}',
-            { query: "London" },
-            { result: { source: "history" } },
-          ),
-        ]),
-      });
-    });
-
-    await act(async () => {});
-    expect(streamCall).not.toHaveBeenCalled();
-  });
-
-  it("does not emit a cancellation onResult for pre-resolved tool calls aborted by reset", async () => {
-    const streamCall = vi.fn();
-    const getTools = () => ({
-      weatherSearch: {
-        parameters: { type: "object", properties: {} },
-        execute: vi.fn(async () => ({ forecast: "ok" })),
-        streamCall,
-      } satisfies Tool,
-    });
-    const onResult = vi.fn();
-    const setToolStatuses = vi.fn();
-
-    const { result, rerender } = renderHook(
-      ({ state }: { state: AssistantTransportState }) =>
-        useToolInvocations({
-          state,
-          getTools,
-          onResult,
-          setToolStatuses,
-        }),
-      {
-        initialProps: {
-          state: createState([]),
-        },
-      },
-    );
-
-    act(() => {
-      rerender({
-        state: createState([
-          createAssistantMessage(
-            '{"query":"London"}',
-            { query: "London" },
-            { result: { source: "history" } },
-          ),
-        ]),
-      });
-    });
-
-    await waitFor(() => {
-      expect(streamCall).toHaveBeenCalledTimes(1);
-    });
-
-    act(() => {
-      result.current.reset();
-    });
-
-    // Flush microtasks through the executor's abort race + the stream
-    // pipeline so any cancellation `result` chunk has a chance to land
-    // before we assert it didn't.
-    for (let i = 0; i < 5; i++) {
-      await act(async () => {});
-    }
-
-    expect(onResult).not.toHaveBeenCalled();
-  });
-});