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

package/src/runtime/useChatRuntime.integration.test.tsx

@@ -0,0 +1,661 @@
+// @vitest-environment jsdom
+
+import { useInsertionEffect, useTapHost } from "@ulearning/ai-tap";
+import { useRuntimeResource } from "@ulearning/ai-core/vue/internal";
+import { DynamicChatTransport } from "./DynamicChatTransport";
+import { useChatThread } from "./useChatThread";
+import { getEventListeners } from "node:events";
+import { act, render, waitFor } from "./__tests__/vue";
+import { screen } from "@testing-library/dom";
+import {
+  AssistantRuntimeProvider,
+  RuntimeAdapterProvider,
+} from "@ulearning/ai-core/vue";
+import { AuiConfig, AuiProvider, useAuiState } from "@ulearning/ai-store";
+import { useAssistantClientDestroySignal } from "@ulearning/ai-store/internal";
+import type { AssistantRuntime } from "@ulearning/ai-core";
+import { AISDKChat } from "./AISDKChat";
+import type { ChatTransport, UIMessage } from "ai";
+import {
+  defineComponent,
+  h,
+  KeepAlive,
+  nextTick,
+  onMounted,
+  shallowRef,
+  Suspense,
+  type VNodeChild,
+} from "vue";
+import { useIsomorphicLayoutEffect } from "@ulearning/ai-store/internal";
+import { describe, expect, it, vi } from "vitest";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import {
+  createCancellableTransport,
+  createStreamHarness,
+  nextTask,
+} from "./__tests__/controlled-transport";
+import { useChatRuntime } from "./useChatRuntime";
+import { useThreadTokenUsage } from "../usage";
+
+const messages: UIMessage[] = [
+  {
+    id: "initial-user-message",
+    role: "user",
+    parts: [{ type: "text", text: "Hello from the server" }],
+  },
+];
+
+const MessageProbe = defineComponent({
+  setup() {
+    const count = useAuiState((state) => state.thread.messages.length);
+    const text = useAuiState(
+      (state) =>
+        state.thread.messages[0]?.parts
+          .map((part) => (part.type === "text" ? part.text : ""))
+          .join("") ?? "",
+    );
+    return () => [
+      h("output", { "data-testid": "message-count" }, count.value),
+      h("output", { "data-testid": "message-text" }, text.value),
+    ];
+  },
+});
+
+const TestApp = defineComponent({
+  setup() {
+    const transport = new AssistantChatTransport({ api: "/api/chat" });
+    const runtime = useChatRuntime({ messages, transport });
+    return () =>
+      h(
+        AssistantRuntimeProvider,
+        { runtime },
+        { default: () => h(MessageProbe) },
+      );
+  },
+});
+
+describe("useChatRuntime integration", () => {
+  it("exposes seeded messages through the mounted thread scope", async () => {
+    render(h(TestApp));
+
+    await waitFor(() => {
+      expect(screen.getByTestId("message-count").textContent).toBe("1");
+      expect(screen.getByTestId("message-text").textContent).toBe(
+        "Hello from the server",
+      );
+    });
+  });
+
+  it("aborts a deleted thread's stream while the host stays mounted", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning, client } = createStreamHarness();
+
+    const view = render(h(StreamingApp, { transport, probe: h(Probe) }));
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+
+    // Wait for the delete after the action has published its state change.
+    let deletion: unknown;
+    await act(async () => {
+      deletion = client().threadListItem.delete();
+    });
+    await deletion;
+
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+    view.unmount();
+  });
+
+  it("aborts the in-flight transport after a real unmount", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
+
+    const view = render(h(StreamingApp, { transport, probe: h(Probe) }));
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+    // Retained-host cycles are covered below; a mounted host must not cancel.
+    expect(getCancelCount()).toBe(0);
+
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+
+  // Real KeepAlive deactivation retains the client and transport.
+  it("keeps streaming while hidden and aborts when the hidden host unmounts", async () => {
+    const { transport, getCancelCount } = createCancellableTransport();
+    const { Probe, send, isRunning } = createStreamHarness();
+
+    let setMode: ((mode: "visible" | "hidden") => void) | undefined;
+    const Shell = defineComponent({
+      setup() {
+        const mode = shallowRef<"visible" | "hidden">("visible");
+        setMode = (value) => {
+          mode.value = value;
+        };
+        return () =>
+          h(KeepAlive, null, {
+            default: () =>
+              mode.value === "visible"
+                ? h(StreamingApp, { transport, probe: h(Probe) })
+                : null,
+          });
+      },
+    });
+    const view = render(h(Shell));
+
+    await act(async () => send());
+    await waitFor(() => expect(isRunning()).toBe(true));
+
+    await act(async () => setMode?.("hidden"));
+    await act(nextTask);
+    expect(getCancelCount()).toBe(0);
+    expect(isRunning()).toBe(true);
+
+    await act(async () => setMode?.("visible"));
+    await act(nextTask);
+    expect(getCancelCount()).toBe(0);
+    expect(isRunning()).toBe(true);
+
+    await act(async () => setMode?.("hidden"));
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+
+  it("aborts a nested runtime's stream when the provider above it unmounts", async () => {
+    const outer = createCancellableTransport();
+    const { transport, getCancelCount } = createCancellableTransport();
+    let nested: AssistantRuntime | undefined;
+
+    // The nested native owner must release its own transport independently.
+    const NestedChat = defineComponent({
+      setup() {
+        nested = useChatRuntime({ transport });
+        return () => null;
+      },
+    });
+    const view = render(
+      h(
+        AuiProvider,
+        {
+          config: AuiConfig({
+            threads: AISDKChat({ transport: outer.transport }),
+          }),
+        },
+        { default: () => h(NestedChat) },
+      ),
+    );
+
+    await waitFor(() => expect(nested).toBeDefined());
+    await act(async () => {
+      await nested!.thread.append("keep streaming");
+    });
+    await waitFor(() => expect(nested!.thread.getState().isRunning).toBe(true));
+
+    view.unmount();
+    await waitFor(() => expect(getCancelCount()).toBe(1));
+  });
+
+  it("aborts a nested runtime when only its own component unmounts", async () => {
+    const outer = createCancellableTransport();
+    const { transport, getCancelCount } = createCancellableTransport();
+    let nested: AssistantRuntime | undefined;
+    let providerSignal: AbortSignal | undefined;
+    let setVisible: ((visible: boolean) => void) | undefined;
+
+    const NestedChat = defineComponent({
+      setup() {
+        providerSignal = useAssistantClientDestroySignal();
+        nested = useChatRuntime({ transport });
+        return () => null;
+      },
+    });
+    const Shell = defineComponent({
+      setup() {
+        const visible = shallowRef(true);
+        setVisible = (value) => {
+          visible.value = value;
+        };
+        const config = AuiConfig({
+          threads: AISDKChat({ transport: outer.transport }),
+        });
+        return () =>
+          h(
+            AuiProvider,
+            { config },
+            { default: () => (visible.value ? h(NestedChat) : null) },
+          );
+      },
+    });
+    const listeners = () => getEventListeners(providerSignal!, "abort").length;
+
+    const view = render(h(Shell));
+    await waitFor(() => expect(nested).toBeDefined());
+    const mounted = listeners();
+
+    for (const cycle of [1, 2]) {
+      await act(async () => {
+        await nested!.thread.append(`stream ${cycle}`);
+      });
+      await waitFor(() =>
+        expect(nested!.thread.getState().isRunning).toBe(true),
+      );
+
+      await act(async () => setVisible?.(false));
+      await waitFor(() => expect(getCancelCount()).toBe(cycle));
+      expect(listeners()).toBe(mounted - 1);
+
+      nested = undefined;
+      await act(async () => setVisible?.(true));
+      await waitFor(() => expect(nested).toBeDefined());
+      expect(listeners()).toBe(mounted);
+    }
+
+    view.unmount();
+    await act(nextTask);
+    expect(getCancelCount()).toBe(2);
+  });
+});
+
+const StreamingApp = defineComponent(
+  (props: { transport: ChatTransport<UIMessage>; probe: VNodeChild }) => {
+    const runtime = useChatRuntime(() => ({ transport: props.transport }));
+    return () =>
+      h(AssistantRuntimeProvider, { runtime }, { default: () => props.probe });
+  },
+  { props: ["transport", "probe"] },
+);
+
+const UsageProbe = defineComponent({
+  setup() {
+    const usage = useThreadTokenUsage();
+    return () =>
+      h(
+        "output",
+        { "data-testid": "total-tokens" },
+        usage.value?.totalTokens ?? "none",
+      );
+  },
+});
+
+const UsageApp = defineComponent({
+  setup() {
+    const transport = new AssistantChatTransport({ api: "/api/chat" });
+    const runtime = useChatRuntime({
+      messages: [
+        ...messages,
+        {
+          id: "assistant-with-usage",
+          role: "assistant",
+          parts: [{ type: "text", text: "Hi" }],
+          metadata: { usage: { inputTokens: 40, outputTokens: 2 } },
+        },
+      ],
+      transport,
+    });
+    return () =>
+      h(
+        AssistantRuntimeProvider,
+        { runtime },
+        { default: () => h(UsageProbe) },
+      );
+  },
+});
+
+describe("useThreadTokenUsage through useChatRuntime", () => {
+  it("reads usage from the message metadata a server attached", async () => {
+    render(h(UsageApp));
+
+    await waitFor(() => {
+      expect(screen.getByTestId("total-tokens").textContent).toBe("42");
+    });
+  });
+});
+
+describe("replacement transports", () => {
+  it("does not loop when an inline transport recreates resumable storage", async () => {
+    const fetch = vi.fn(
+      async () =>
+        new Response(
+          new ReadableStream({ start: (controller) => controller.close() }),
+          { headers: { "content-type": "text/event-stream" } },
+        ),
+    );
+    let renders = 0;
+    let sent = false;
+    const SendOnLayout = defineComponent(
+      (props: { runtime: AssistantRuntime }) => {
+        useIsomorphicLayoutEffect(
+          () => {
+            if (sent) return;
+            sent = true;
+            void props.runtime.thread.append("hello");
+          },
+          () => [props.runtime],
+        );
+        return () => null;
+      },
+      { props: ["runtime"] },
+    );
+    const NestedChat = defineComponent({
+      setup() {
+        const runtime = useChatRuntime(() => {
+          renders += 1;
+          return {
+            transport: new AssistantChatTransport({
+              fetch,
+              resumable: {
+                storage: {
+                  getStreamId: () => null,
+                  setStreamId: vi.fn(),
+                  clear: vi.fn(),
+                },
+                resumeApi: "/api/chat/resume",
+              },
+            }),
+          };
+        });
+        return () => h(SendOnLayout, { runtime });
+      },
+    });
+    render(
+      h(
+        AuiProvider,
+        {
+          config: AuiConfig({
+            threads: AISDKChat({ transport: new AssistantChatTransport() }),
+          }),
+        },
+        { default: () => h(NestedChat) },
+      ),
+    );
+    await waitFor(() => expect(fetch).toHaveBeenCalledOnce());
+    expect(renders).toBeLessThan(10);
+  });
+
+  it("routes sends through a replacement transport", async () => {
+    const createEmptyStream = () =>
+      new ReadableStream({ start: (controller) => controller.close() });
+    const sendA = vi.fn(async () => createEmptyStream());
+    const sendB = vi.fn(async () => createEmptyStream());
+    const transportA: ChatTransport<UIMessage> = {
+      sendMessages: sendA,
+      reconnectToStream: vi.fn(),
+    };
+    const transportB: ChatTransport<UIMessage> = {
+      sendMessages: sendB,
+      reconnectToStream: vi.fn(),
+    };
+    const SendOnLayout = defineComponent(
+      (props: { runtime: AssistantRuntime }) => {
+        useIsomorphicLayoutEffect(
+          () => {
+            void props.runtime.thread.append({
+              role: "user",
+              content: [{ type: "text", text: "hello" }],
+            });
+          },
+          () => [props.runtime],
+        );
+        return () => null;
+      },
+      { props: ["runtime"] },
+    );
+    const App = defineComponent(
+      (props: { transport: ChatTransport<UIMessage>; send?: boolean }) => {
+        const runtime = useChatRuntime(() => ({ transport: props.transport }));
+        return () =>
+          h(
+            AssistantRuntimeProvider,
+            { runtime },
+            {
+              default: () => (props.send ? h(SendOnLayout, { runtime }) : null),
+            },
+          );
+      },
+      { props: ["transport", "send"] },
+    );
+    const view = render(h(App, { transport: transportA }));
+    await view.rerender(h(App, { transport: transportB, send: true }));
+    await waitFor(() => expect(sendB).toHaveBeenCalledOnce());
+    expect(sendA).not.toHaveBeenCalled();
+  });
+
+  it("routes sends triggered during the initial layout commit", async () => {
+    const send = vi.fn(
+      async () =>
+        new ReadableStream({ start: (controller) => controller.close() }),
+    );
+    const transport: ChatTransport<UIMessage> = {
+      sendMessages: send,
+      reconnectToStream: vi.fn(),
+    };
+    let resolveLayout!: () => void;
+    const layoutCommitted = new Promise<void>((resolve) => {
+      resolveLayout = resolve;
+    });
+    const SendOnLayout = defineComponent(
+      (props: { runtime: AssistantRuntime }) => {
+        onMounted(() => {
+          void props.runtime.thread.append({
+            role: "user",
+            content: [{ type: "text", text: "hello" }],
+          });
+          resolveLayout();
+        });
+        return () => null;
+      },
+      { props: ["runtime"] },
+    );
+    const App = defineComponent({
+      setup() {
+        const runtime = useChatRuntime({ transport });
+        return () =>
+          h(
+            AssistantRuntimeProvider,
+            { runtime },
+            { default: () => h(SendOnLayout, { runtime }) },
+          );
+      },
+    });
+    const view = render(h(App));
+    try {
+      await layoutCommitted;
+      await waitFor(() => expect(send).toHaveBeenCalledOnce());
+    } finally {
+      view.unmount();
+    }
+  });
+
+  it("keeps AssistantChatTransport wired through native retained-scope effect replay", async () => {
+    const sourceTransport = new AssistantChatTransport<UIMessage>({
+      api: "/api/chat",
+    });
+    const transport = new DynamicChatTransport(sourceTransport);
+    const wiredDuringEffectSetup: boolean[] = [];
+    const Probe = defineComponent(
+      (props: { effects: () => void }) => {
+        useIsomorphicLayoutEffect(() => {
+          props.effects();
+          try {
+            wiredDuringEffectSetup.push(
+              transport.getCurrentTransport("strict-mode-thread") !==
+                sourceTransport,
+            );
+          } catch {
+            wiredDuringEffectSetup.push(false);
+          }
+        }, []);
+        return () => null;
+      },
+      { props: ["effects"] },
+    );
+    const App = defineComponent({
+      setup() {
+        const host = useTapHost(function ChatThreadResource() {
+          return useChatThread(
+            { transport },
+            {
+              id: "strict-mode-thread",
+              isMainThread: true,
+              getThreadListItem: () => undefined,
+            },
+          );
+        });
+        return () => h(Probe, { effects: host.effects });
+      },
+    });
+    const visible = shallowRef(true);
+    render(
+      h(
+        defineComponent({
+          setup: () => () =>
+            h(KeepAlive, null, {
+              default: () => (visible.value ? h(App) : null),
+            }),
+        }),
+      ),
+    );
+    visible.value = false;
+    await nextTick();
+    visible.value = true;
+    await nextTick();
+    expect(wiredDuringEffectSetup).toEqual([true, true]);
+  });
+
+  it("does not publish transport contexts from suspended renders", () => {
+    const sourceTransport = new AssistantChatTransport<UIMessage>({
+      api: "/api/chat",
+    });
+    const transport = new DynamicChatTransport(sourceTransport);
+    const pending = new Promise<never>(() => {});
+    const Blocker = defineComponent({
+      async setup() {
+        await pending;
+        return () => null;
+      },
+    });
+    const SuspendedThread = defineComponent({
+      setup() {
+        useTapHost(function ChatThreadResource() {
+          return useChatThread(
+            { transport },
+            {
+              id: "suspended-thread",
+              isMainThread: true,
+              getThreadListItem: () => undefined,
+            },
+          );
+        });
+        return () => h(Blocker);
+      },
+    });
+    render(
+      h(Suspense, null, {
+        default: () => h(SuspendedThread),
+        fallback: () => null,
+      }),
+    );
+
+    expect(transport.getCurrentTransport("suspended-thread")).toBe(
+      sourceTransport,
+    );
+  });
+
+  it("routes runtime sends through a wired clone with the latest thread item", async () => {
+    const bodies: Array<{ id: string; system: string }> = [];
+    const completeRequests: (() => void)[] = [];
+    const sourceTransport = new AssistantChatTransport<UIMessage>({
+      fetch: vi.fn(async (_input, init) => {
+        bodies.push(JSON.parse(String(init?.body)));
+        return new Response(
+          new ReadableStream({
+            start(controller) {
+              completeRequests.push(() => controller.close());
+            },
+          }),
+          { headers: { "content-type": "text/event-stream" } },
+        );
+      }),
+    });
+    const transport = new DynamicChatTransport(sourceTransport);
+    const remoteId = shallowRef("remote-a");
+    let publishedRemoteId: string | undefined;
+    let runtime!: AssistantRuntime;
+    const Probe = defineComponent({
+      setup() {
+        runtime = useRuntimeResource(() => {
+          const committedRemoteId = remoteId.value;
+          const current = useChatThread(
+            { transport },
+            {
+              id: "stable-thread",
+              isMainThread: true,
+              getThreadListItem: () => ({
+                initialize: async () => ({
+                  remoteId: committedRemoteId,
+                  externalId: undefined,
+                }),
+              }),
+            },
+          );
+          useInsertionEffect(() => {
+            publishedRemoteId = committedRemoteId;
+          });
+          return current;
+        });
+        return () => null;
+      },
+    });
+    render(
+      h(
+        RuntimeAdapterProvider,
+        {
+          adapters: {
+            modelContext: {
+              getModelContext: () => ({ system: "system prompt" }),
+            },
+          },
+        },
+        { default: () => h(Probe) },
+      ),
+    );
+    const result = {
+      get current() {
+        return runtime;
+      },
+    };
+    const rerender = async (props: { remoteId: string }) => {
+      await act(() => {
+        remoteId.value = props.remoteId;
+      });
+      await waitFor(() => expect(publishedRemoteId).toBe(props.remoteId));
+    };
+    const send = async () => {
+      const index = bodies.length;
+      await act(async () => {
+        await result.current.thread.append({
+          role: "user",
+          content: [{ type: "text", text: "hello" }],
+        });
+      });
+      await waitFor(() => {
+        expect(bodies).toHaveLength(index + 1);
+        expect(result.current.thread.getState().isRunning).toBe(true);
+      });
+      completeRequests[index]!();
+      await waitFor(() =>
+        expect(result.current.thread.getState().isRunning).toBe(false),
+      );
+    };
+
+    await send();
+    await rerender({ remoteId: "remote-b" });
+    await send();
+
+    expect(bodies).toEqual([
+      expect.objectContaining({ id: "remote-a", system: "system prompt" }),
+      expect.objectContaining({ id: "remote-b", system: "system prompt" }),
+    ]);
+  });
+});