@ulearning/ai-sdk-adapter 0.0.0-stage → 1.0.0-next.0

package/src/runtime/useAISDKRuntime.queue.test.tsx

@@ -0,0 +1,707 @@
+// @vitest-environment jsdom
+
+import {
+  act,
+  cleanup,
+  renderHook,
+  useTestChat,
+  waitFor,
+} from "./__tests__/tap";
+import { render } from "./__tests__/vue";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { defineComponent, h, onMounted } from "vue";
+import { Chat } from "@ai-sdk/vue";
+import { useRuntimeResource } from "@ulearning/ai-core/vue/internal";
+import { useChatSnapshot } from "./useChatSnapshot";
+const screen = {
+  getByTestId(id: string) {
+    const element = document.querySelector(`[data-testid="${id}"]`);
+    if (!element) throw new Error(`Missing test element: ${id}`);
+    return element;
+  },
+};
+import {
+  lastAssistantMessageIsCompleteWithToolCalls,
+  type ChatTransport,
+  type UIMessage,
+  type UIMessageChunk,
+} from "ai";
+import type { AppendMessage, AssistantRuntime } from "@ulearning/ai-core";
+import { AssistantRuntimeProvider } from "@ulearning/ai-core/vue";
+import { useAuiState } from "@ulearning/ai-store";
+import { useAISDKRuntime, type AISDKRuntimeAdapter } from "./useAISDKRuntime";
+import { useAISDKError } from "../hooks";
+import { toCreateMessage } from "../converters/toCreateMessage";
+
+const textOf = (message: UIMessage | undefined) =>
+  message?.parts
+    .map((part) => (part.type === "text" ? part.text : ""))
+    .join("") ?? "";
+
+const createTransport = () => {
+  const events: string[] = [];
+  const requests: {
+    prompt: string;
+    messages: string[];
+    rawMessages: UIMessage[];
+    signal: AbortSignal | undefined;
+    stream: (text: string, messageId?: string) => void;
+    emit: (...chunks: UIMessageChunk[]) => void;
+    finish: () => void;
+  }[] = [];
+  let held: Promise<void> | undefined;
+  let release = () => {};
+  const transport: ChatTransport<UIMessage> = {
+    async sendMessages({ messages, abortSignal }) {
+      const prompt = textOf(messages.findLast((m) => m.role === "user"));
+      events.push(`send ${prompt}`);
+      let controller!: ReadableStreamDefaultController<UIMessageChunk>;
+      const stream = new ReadableStream<UIMessageChunk>({
+        start: (c) => {
+          controller = c;
+        },
+      });
+      requests.push({
+        prompt,
+        messages: messages.map(textOf).filter(Boolean),
+        rawMessages: messages,
+        signal: abortSignal,
+        stream: (text, messageId) => {
+          controller.enqueue({
+            type: "start",
+            ...(messageId && { messageId }),
+          });
+          controller.enqueue({ type: "text-start", id: "t" });
+          controller.enqueue({ type: "text-delta", id: "t", delta: text });
+        },
+        emit: (...chunks) => {
+          for (const chunk of chunks) controller.enqueue(chunk);
+        },
+        finish: () => {
+          controller.enqueue({ type: "text-end", id: "t" });
+          controller.enqueue({ type: "finish" });
+          controller.close();
+        },
+      });
+      await held;
+      return stream;
+    },
+    reconnectToStream: async () => null,
+  };
+  return {
+    transport,
+    events,
+    requests,
+    hold: () => {
+      held = new Promise((resolve) => {
+        release = () => {
+          held = undefined;
+          resolve();
+        };
+      });
+    },
+    release: () => release(),
+  };
+};
+
+const QueuedPrompts = defineComponent({
+  setup() {
+    const prompts = useAuiState((s) =>
+      s.composer.queue.map((item) => item.prompt).join(","),
+    );
+    return () => h("output", { "data-testid": "queued" }, prompts.value);
+  },
+});
+const QueueError = defineComponent({
+  setup() {
+    const error = useAISDKError();
+    return () =>
+      h("output", { "data-testid": "queue-error" }, error.value?.message);
+  },
+});
+
+const setup = (
+  options: AISDKRuntimeAdapter = {
+    unstable_enableMessageQueue: true,
+  },
+) => {
+  const harness = createTransport();
+  const { result, rerender } = renderHook(
+    (runtimeOptions) => {
+      const chat = useTestChat({
+        transport: harness.transport,
+        onFinish: ({ message, isAbort }) => {
+          harness.events.push(
+            `finish ${isAbort ? "aborted" : textOf(message)}`,
+          );
+        },
+      });
+      return useAISDKRuntime(chat, runtimeOptions);
+    },
+    { initialProps: options },
+  );
+  const runtime = () => result.current as AssistantRuntime;
+  render(
+    h(
+      AssistantRuntimeProvider,
+      { runtime: runtime() },
+      { default: () => [h(QueuedPrompts), h(QueueError)] },
+    ),
+  );
+  const send = (text: string, steer?: boolean) =>
+    act(() => {
+      runtime().thread.composer.setText(text);
+      runtime().thread.composer.send(steer === undefined ? {} : { steer });
+    });
+  const queued = () => screen.getByTestId("queued").textContent;
+  const isRunning = () => runtime().thread.getState().isRunning;
+  return { ...harness, runtime, send, queued, isRunning, setOptions: rerender };
+};
+
+afterEach(() => {
+  cleanup();
+});
+
+describe("useAISDKRuntime unstable_enableMessageQueue", () => {
+  it("does not let a disabled queue's error hide a later direct-send error", async () => {
+    let fail = true;
+    function createMessage<UI_MESSAGE extends UIMessage>(
+      message: AppendMessage,
+    ) {
+      if (fail) throw new Error("queue conversion failed");
+      return toCreateMessage<UI_MESSAGE>(message);
+    }
+    const options = {
+      unstable_enableMessageQueue: true,
+      toCreateMessage: createMessage,
+    };
+    const logged = vi.spyOn(console, "error").mockImplementation(() => {});
+    try {
+      const { send, requests, setOptions } = setup(options);
+      await send("first");
+      await waitFor(() =>
+        expect(screen.getByTestId("queue-error").textContent).toBe(
+          "queue conversion failed",
+        ),
+      );
+      await act(async () =>
+        setOptions({ ...options, unstable_enableMessageQueue: false }),
+      );
+      expect(screen.getByTestId("queue-error").textContent).toBe("");
+
+      fail = false;
+      await send("second");
+      await waitFor(() => expect(requests).toHaveLength(1));
+      act(() =>
+        requests[0]!.emit({
+          type: "error",
+          errorText: "direct request failed",
+        }),
+      );
+      await waitFor(() =>
+        expect(screen.getByTestId("queue-error").textContent).toBe(
+          "direct request failed",
+        ),
+      );
+      await act(async () => setOptions(options));
+      expect(screen.getByTestId("queue-error").textContent).toBe(
+        "direct request failed",
+      );
+    } finally {
+      logged.mockRestore();
+    }
+  });
+
+  it("does not mark an earlier answer failed when a queued prompt fails before append", async () => {
+    let fail = false;
+    function createMessage<UI_MESSAGE extends UIMessage>(
+      message: AppendMessage,
+    ) {
+      if (fail) throw new Error("conversion failed");
+      return toCreateMessage<UI_MESSAGE>(message);
+    }
+    const logged = vi.spyOn(console, "error").mockImplementation(() => {});
+    try {
+      const { send, requests, runtime, isRunning, queued } = setup({
+        unstable_enableMessageQueue: true,
+        toCreateMessage: createMessage,
+      });
+      await send("first");
+      await waitFor(() => expect(requests).toHaveLength(1));
+      await act(async () => {
+        requests[0]!.stream("answer");
+        requests[0]!.finish();
+      });
+      await waitFor(() => expect(isRunning()).toBe(false));
+      const status = runtime().thread.getState().messages[1]!.status;
+      fail = true;
+      await send("second");
+      await waitFor(() =>
+        expect(screen.getByTestId("queue-error").textContent).toBe(
+          "conversion failed",
+        ),
+      );
+      expect(queued()).toBe("second");
+      expect(runtime().thread.getState().messages).toHaveLength(2);
+      expect(runtime().thread.getState().messages[1]!.status).toEqual(status);
+    } finally {
+      logged.mockRestore();
+    }
+  });
+
+  it("retains a failed prompt and exposes the error until an explicit retry", async () => {
+    const error = new Error("message conversion failed");
+    let fail = true;
+    function createMessage<UI_MESSAGE extends UIMessage>(
+      message: AppendMessage,
+    ) {
+      if (fail) throw error;
+      return toCreateMessage<UI_MESSAGE>(message);
+    }
+    const logged = vi.spyOn(console, "error").mockImplementation(() => {});
+    try {
+      const { send, requests, queued, isRunning } = setup({
+        unstable_enableMessageQueue: true,
+        toCreateMessage: createMessage,
+      });
+      await send("first");
+      await waitFor(() =>
+        expect(screen.getByTestId("queue-error").textContent).toBe(
+          error.message,
+        ),
+      );
+      expect(queued()).toBe("first");
+      expect(requests).toHaveLength(0);
+
+      fail = false;
+      await send("second", false);
+      await waitFor(() => expect(requests).toHaveLength(1));
+      expect(requests[0]!.prompt).toBe("first");
+      expect(queued()).toBe("second");
+      expect(screen.getByTestId("queue-error").textContent).toBe("");
+      await act(async () => {
+        requests[0]!.stream("one");
+        requests[0]!.finish();
+      });
+      await waitFor(() => expect(requests).toHaveLength(2));
+      expect(requests[1]!.prompt).toBe("second");
+      await act(async () => {
+        requests[1]!.stream("two");
+        requests[1]!.finish();
+      });
+      await waitFor(() => expect(isRunning()).toBe(false));
+      expect(logged).toHaveBeenCalled();
+    } finally {
+      logged.mockRestore();
+    }
+  });
+
+  it("holds sends made during a run and sends them one request at a time", async () => {
+    const { requests, events, send, queued, isRunning, runtime } = setup();
+
+    await send("first");
+    await waitFor(() => expect(requests).toHaveLength(1));
+    expect(runtime().thread.getState().capabilities.queue).toBe(true);
+
+    await send("second", false);
+    await send("third", false);
+    await waitFor(() => expect(queued()).toBe("second,third"));
+    expect(requests).toHaveLength(1);
+
+    await act(async () => {
+      requests[0]!.stream("one");
+      requests[0]!.finish();
+    });
+    await waitFor(() => expect(requests).toHaveLength(2));
+    expect(requests[1]!.prompt).toBe("second");
+    await waitFor(() => expect(queued()).toBe("third"));
+
+    await act(async () => {
+      requests[1]!.stream("two");
+      requests[1]!.finish();
+    });
+    await waitFor(() => expect(requests).toHaveLength(3));
+    expect(requests[2]!.prompt).toBe("third");
+
+    await act(async () => {
+      requests[2]!.stream("three");
+      requests[2]!.finish();
+    });
+    await waitFor(() => expect(isRunning()).toBe(false));
+    expect(queued()).toBe("");
+    expect(events).toEqual([
+      "send first",
+      "finish one",
+      "send second",
+      "finish two",
+      "send third",
+      "finish three",
+    ]);
+  });
+
+  it("steering waits for the running request, then goes before follow-ups", async () => {
+    const { requests, events, hold, release, send, queued, isRunning } =
+      setup();
+
+    hold();
+    await send("first");
+    await waitFor(() => expect(requests).toHaveLength(1));
+    await waitFor(() => expect(isRunning()).toBe(true));
+
+    await send("third", false);
+    await send("second", true);
+    expect(requests[0]!.signal?.aborted).toBe(false);
+    expect(queued()).toBe("second,third");
+    expect(requests).toHaveLength(1);
+
+    await act(async () => {
+      release();
+    });
+    expect(requests).toHaveLength(1);
+    await act(async () => {
+      requests[0]!.stream("one");
+      requests[0]!.finish();
+    });
+    await waitFor(() => expect(requests).toHaveLength(2));
+    expect(requests[1]!.prompt).toBe("second");
+    expect(requests[1]!.messages).toEqual(["first", "one", "second"]);
+
+    await act(async () => {
+      requests[1]!.stream("two");
+      requests[1]!.finish();
+    });
+    await waitFor(() => expect(requests).toHaveLength(3));
+    expect(requests[2]!.prompt).toBe("third");
+    await act(async () => {
+      requests[2]!.stream("three");
+      requests[2]!.finish();
+    });
+    await waitFor(() => expect(isRunning()).toBe(false));
+    expect(events).toEqual([
+      "send first",
+      "finish one",
+      "send second",
+      "finish two",
+      "send third",
+      "finish three",
+    ]);
+  });
+
+  it("keeps queued messages when the run is cancelled and sends them on the next send", async () => {
+    const { requests, send, queued, isRunning, runtime } = setup();
+
+    await send("first");
+    await waitFor(() => expect(requests).toHaveLength(1));
+    await act(async () => {
+      requests[0]!.stream("partial");
+    });
+    await waitFor(() =>
+      expect(
+        runtime().thread.getState().messages.at(-1)?.content,
+      ).toContainEqual(
+        expect.objectContaining({ type: "text", text: "partial" }),
+      ),
+    );
+    await send("second", false);
+    await waitFor(() => expect(queued()).toBe("second"));
+
+    act(() => {
+      runtime().thread.cancelRun();
+    });
+    await waitFor(() => expect(isRunning()).toBe(false));
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    expect(requests).toHaveLength(1);
+    expect(queued()).toBe("second");
+
+    await send("third", false);
+    await waitFor(() => expect(requests).toHaveLength(2));
+    expect(requests[1]!.prompt).toBe("second");
+    await waitFor(() => expect(queued()).toBe("third"));
+  });
+
+  const cancelWithQueued = async (harness: ReturnType<typeof setup>) => {
+    const { requests, send, queued, isRunning, runtime } = harness;
+    await send("first");
+    await waitFor(() => expect(requests).toHaveLength(1));
+    await act(async () => {
+      requests[0]!.stream("partial");
+    });
+    await waitFor(() =>
+      expect(
+        runtime().thread.getState().messages.at(-1)?.content,
+      ).toContainEqual(
+        expect.objectContaining({ type: "text", text: "partial" }),
+      ),
+    );
+    await send("second", false);
+    await waitFor(() => expect(queued()).toBe("second"));
+    act(() => {
+      runtime().thread.cancelRun();
+    });
+    await waitFor(() => expect(isRunning()).toBe(false));
+    expect(queued()).toBe("second");
+  };
+
+  it("drops queued messages when a message is edited", async () => {
+    const harness = setup();
+    const { requests, queued, isRunning, runtime } = harness;
+    await cancelWithQueued(harness);
+
+    act(() => {
+      const composer = runtime().thread.getMessageByIndex(0).composer;
+      composer.beginEdit();
+      composer.setText("edited");
+      composer.send();
+    });
+    await waitFor(() => expect(requests).toHaveLength(2));
+    expect(requests[1]!.prompt).toBe("edited");
+    expect(queued()).toBe("");
+
+    await act(async () => {
+      requests[1]!.stream("two");
+      requests[1]!.finish();
+    });
+    await waitFor(() => expect(isRunning()).toBe(false));
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    expect(requests).toHaveLength(2);
+  });
+
+  it("drops queued messages when a response is reloaded", async () => {
+    const harness = setup();
+    const { requests, queued, isRunning, runtime } = harness;
+    await cancelWithQueued(harness);
+
+    act(() => {
+      runtime().thread.getMessageByIndex(1).reload();
+    });
+    await waitFor(() => expect(requests).toHaveLength(2));
+    expect(requests[1]!.prompt).toBe("first");
+    expect(queued()).toBe("");
+
+    await act(async () => {
+      requests[1]!.stream("again");
+      requests[1]!.finish();
+    });
+    await waitFor(() => expect(isRunning()).toBe(false));
+    await act(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+    });
+    expect(requests).toHaveLength(2);
+  });
+
+  it("keeps a cancelled steer queued alongside the restored unanswered draft", async () => {
+    const { requests, hold, release, send, queued, isRunning, runtime } =
+      setup();
+
+    hold();
+    await send("first");
+    await waitFor(() => expect(requests).toHaveLength(1));
+    await waitFor(() => expect(isRunning()).toBe(true));
+    await send("second", true);
+
+    act(() => {
+      runtime().thread.cancelRun();
+    });
+    await act(async () => {
+      release();
+    });
+    await waitFor(() => expect(queued()).toBe("second"));
+    expect(requests).toHaveLength(1);
+    await waitFor(() => expect(isRunning()).toBe(false));
+    expect(runtime().thread.getState().messages).toEqual([]);
+    expect(runtime().thread.composer.getState().text).toBe("first");
+    expect(screen.getByTestId("queue-error").textContent).toBe("");
+
+    await send("third");
+    await waitFor(() => expect(requests).toHaveLength(2));
+    expect(requests[1]!.prompt).toBe("second");
+    expect(queued()).toBe("third");
+    await act(async () => {
+      requests[1]!.stream("two");
+      requests[1]!.finish();
+    });
+    await waitFor(() => expect(requests).toHaveLength(3));
+    expect(requests[2]!.prompt).toBe("third");
+    expect(queued()).toBe("");
+  });
+
+  it("lets a client tool finish and sends its result before the steered message", async () => {
+    const harness = createTransport();
+    let shouldContinue = true;
+    let finishTool!: (result: string) => void;
+    const execute = vi.fn(
+      (_args: unknown, _context: { abortSignal: AbortSignal }) =>
+        new Promise<string>((resolve) => {
+          finishTool = resolve;
+        }),
+    );
+    const { result } = renderHook(() =>
+      useAISDKRuntime(
+        useTestChat({
+          transport: harness.transport,
+          sendAutomaticallyWhen: (options) => {
+            if (!shouldContinue) return false;
+            const complete =
+              lastAssistantMessageIsCompleteWithToolCalls(options);
+            if (complete) shouldContinue = false;
+            return complete;
+          },
+        }),
+        { unstable_enableMessageQueue: true },
+      ),
+    );
+    result.current.registerModelContextProvider({
+      getModelContext: () => ({
+        tools: {
+          weather: {
+            parameters: { type: "object", properties: {} },
+            execute,
+          },
+        },
+      }),
+    });
+
+    act(() => {
+      result.current.thread.composer.setText("first");
+      result.current.thread.composer.send();
+    });
+    await waitFor(() => expect(harness.requests).toHaveLength(1));
+    await act(async () => {
+      harness.requests[0]!.emit(
+        { type: "start" },
+        {
+          type: "tool-input-available",
+          toolCallId: "tool-1",
+          toolName: "weather",
+          input: {},
+        },
+      );
+      harness.requests[0]!.finish();
+    });
+    await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+    expect(result.current.thread.getState().isRunning).toBe(true);
+
+    act(() => {
+      result.current.thread.composer.setText("second");
+      result.current.thread.composer.send({ steer: true });
+    });
+    expect(harness.requests).toHaveLength(1);
+    expect(execute.mock.calls[0]![1].abortSignal.aborted).toBe(false);
+
+    await act(async () => {
+      finishTool("sunny");
+    });
+    await waitFor(() => expect(harness.requests).toHaveLength(2));
+    expect(harness.requests[1]!.prompt).toBe("first");
+    expect(harness.requests[1]!.messages).not.toContain("second");
+    expect(JSON.stringify(harness.requests[1]!.rawMessages)).toContain("sunny");
+    expect(JSON.stringify(harness.requests[1]!.rawMessages)).toContain(
+      "tool-1",
+    );
+    expect(execute.mock.calls[0]![1].abortSignal.aborted).toBe(false);
+
+    await act(async () => {
+      harness.requests[1]!.stream("continued", "continuation");
+      harness.requests[1]!.finish();
+    });
+    await waitFor(() => expect(harness.requests).toHaveLength(3));
+    expect(harness.requests[2]!.prompt).toBe("second");
+    expect(execute.mock.calls[0]![1].abortSignal.aborted).toBe(false);
+    await act(async () => {
+      harness.requests[2]!.stream("two");
+      harness.requests[2]!.finish();
+    });
+    await waitFor(() =>
+      expect(result.current.thread.getState().isRunning).toBe(false),
+    );
+  });
+
+  it("holds queued messages while sending is disabled", async () => {
+    const harness = createTransport();
+    const { result, rerender } = renderHook(
+      ({ isSendDisabled }: { isSendDisabled: boolean }) =>
+        useAISDKRuntime(useTestChat({ transport: harness.transport }), {
+          unstable_enableMessageQueue: true,
+          isSendDisabled,
+        }),
+      { initialProps: { isSendDisabled: false } },
+    );
+
+    act(() => {
+      result.current.thread.append("first");
+    });
+    await waitFor(() => expect(harness.requests).toHaveLength(1));
+    act(() => {
+      result.current.thread.composer.setText("second");
+      result.current.thread.composer.send({ steer: false });
+    });
+    rerender({ isSendDisabled: true });
+
+    await act(async () => {
+      harness.requests[0]!.stream("one");
+      harness.requests[0]!.finish();
+    });
+    await waitFor(() =>
+      expect(result.current.thread.getState().isRunning).toBe(false),
+    );
+    expect(harness.requests).toHaveLength(1);
+
+    rerender({ isSendDisabled: false });
+    await waitFor(() => expect(harness.requests).toHaveLength(2));
+    expect(harness.requests[1]!.prompt).toBe("second");
+  });
+
+  it("holds a send made before its first effect when it mounts over a running chat", async () => {
+    const harness = createTransport();
+    const chat = new Chat({ transport: harness.transport });
+    void chat.sendMessage({ text: "first" });
+    await waitFor(() => expect(harness.requests).toHaveLength(1));
+    expect(["submitted", "streaming"]).toContain(chat.status);
+
+    const SendOnMount = defineComponent<{ runtime: AssistantRuntime }>(
+      (props) => {
+        onMounted(() => {
+          props.runtime.thread.composer.setText("second");
+          props.runtime.thread.composer.send({ steer: false });
+        });
+        return () => null;
+      },
+      { props: ["runtime"] },
+    );
+    const Host = defineComponent({
+      setup() {
+        const runtime = useRuntimeResource(() =>
+          useAISDKRuntime(useChatSnapshot(chat, undefined), {
+            unstable_enableMessageQueue: true,
+          }),
+        );
+        return () => h(SendOnMount, { runtime });
+      },
+    });
+    const root = render(h(Host));
+    try {
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(harness.requests).toHaveLength(1);
+
+      harness.requests[0]!.stream("one");
+      harness.requests[0]!.finish();
+      await waitFor(() => expect(harness.requests).toHaveLength(2));
+      expect(harness.requests[1]!.prompt).toBe("second");
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("exposes no queue when the option is omitted", async () => {
+    const { requests, runtime, send, queued } = setup({});
+
+    expect(runtime().thread.getState().capabilities.queue).toBe(false);
+    await send("first");
+    await waitFor(() => expect(requests).toHaveLength(1));
+    expect(queued()).toBe("");
+  });
+});