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

package/src/runtime/AISDKThreads.test.ts

@@ -0,0 +1,811 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { lastAssistantMessageIsCompleteWithToolCalls } from "ai";
+
+import { flushTapSync } from "@ulearning/ai-tap";
+import { AuiConfig, createAssistantClient } from "@ulearning/ai-store/client";
+import { AISDKThreads } from "./AISDKThreads";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import {
+  createCancellableTransport,
+  createControlledTransport,
+} from "./__tests__/controlled-transport";
+
+const waitFor = (assertion: () => void) =>
+  vi.waitFor(assertion, { interval: 10 });
+
+const textReply = (text: string) =>
+  [
+    { type: "start" },
+    { type: "text-start", id: "t1" },
+    { type: "text-delta", id: "t1", delta: text },
+    { type: "text-end", id: "t1" },
+    { type: "finish" },
+  ] as never[];
+
+const threadText = (aui: ReturnType<typeof createAssistantClient>) =>
+  aui
+    .getClient()
+    .thread.getState()
+    .messages.map((m) =>
+      m.content.map((part) => (part.type === "text" ? part.text : "")).join(""),
+    );
+
+const createLiveHandle = (
+  getOptions: () => Parameters<typeof AISDKThreads>[0],
+) => {
+  const listeners = new Set<() => void>();
+  const handle = createAssistantClient({
+    getConfig: () => AuiConfig({ threads: AISDKThreads(getOptions()) }),
+    subscribe: (listener) => {
+      listeners.add(listener);
+      return () => listeners.delete(listener);
+    },
+  });
+  handle.subscribe(() => {});
+  return {
+    handle,
+    rerender: () =>
+      flushTapSync(() => listeners.forEach((listener) => listener())),
+  };
+};
+
+const chatCallbacks = () => ({
+  onToolCall: vi.fn(),
+  onData: vi.fn(),
+  onFinish: vi.fn(),
+  onError: vi.fn(),
+  sendAutomaticallyWhen: vi.fn(() => false),
+});
+
+describe("AISDKThreads", () => {
+  const startTool = async (
+    execute: (
+      args: unknown,
+      context: {
+        abortSignal: AbortSignal;
+        human: (payload: unknown) => Promise<unknown>;
+      },
+    ) => Promise<unknown>,
+  ) => {
+    const { transport, emit, close } = createControlledTransport();
+    const sendMessages = vi.spyOn(transport, "sendMessages");
+    const handle = createAssistantClient(
+      AuiConfig({
+        threads: AISDKThreads({
+          transport,
+          sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+        }),
+      }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+    flushTapSync(() =>
+      aui.modelContext.register({
+        getModelContext: () => ({
+          tools: {
+            lookup: {
+              parameters: { type: "object", properties: {} },
+              execute,
+            },
+          },
+        }),
+      }),
+    );
+    flushTapSync(() => aui.composer.setText("look up"));
+    flushTapSync(() => aui.composer.send());
+    await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(1));
+    emit(
+      { type: "start", messageId: "assistant-1" },
+      {
+        type: "tool-input-available",
+        toolCallId: "call-1",
+        toolName: "lookup",
+        input: {},
+      },
+      { type: "finish" },
+    );
+    close();
+    return { handle, aui, sendMessages };
+  };
+
+  it.each([
+    ["pending", (_signal: AbortSignal) => new Promise<string>(() => {})],
+    [
+      "rejecting with AbortError",
+      (signal: AbortSignal) =>
+        new Promise<string>((_resolve, reject) => {
+          signal.addEventListener(
+            "abort",
+            () => reject(new DOMException("aborted", "AbortError")),
+            { once: true },
+          );
+        }),
+    ],
+  ])(
+    "keeps a stopped frontend tool result from starting another request when $0",
+    async (_settlement, createExecution) => {
+      let signal: AbortSignal | undefined;
+      const execute = vi.fn(
+        (_args: unknown, context: { abortSignal: AbortSignal }) => {
+          signal = context.abortSignal;
+          return createExecution(signal);
+        },
+      );
+      const { handle, aui, sendMessages } = await startTool(execute);
+      try {
+        await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+        await waitFor(() => expect(aui.thread.getState().isRunning).toBe(true));
+
+        flushTapSync(() => aui.thread.cancelRun());
+
+        await waitFor(() => expect(signal?.aborted).toBe(true));
+        await waitFor(() => {
+          const toolCall = aui.thread
+            .getState()
+            .messages.filter((message) => message.role === "assistant")
+            .flatMap((message) => message.content)
+            .find(
+              (part) =>
+                part.type === "tool-call" && part.toolCallId === "call-1",
+            );
+          expect(toolCall?.type === "tool-call" && toolCall.isError).toBe(true);
+        });
+        await new Promise((resolve) => setTimeout(resolve, 20));
+        expect(sendMessages).toHaveBeenCalledTimes(1);
+      } finally {
+        handle.destroy();
+      }
+    },
+  );
+
+  it("still auto-sends after a successful frontend tool result", async () => {
+    const execute = vi.fn(async () => "done");
+    const { handle, sendMessages } = await startTool(execute);
+    try {
+      await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(2));
+    } finally {
+      handle.destroy();
+    }
+  });
+
+  it("emits the selected thread id when its composer sends", async () => {
+    const { transport } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+      const sent = vi.fn();
+      aui.on({ scope: "thread", event: "composer.send" }, sent);
+
+      flushTapSync(() => aui.composer.setText("hello"));
+      flushTapSync(() => aui.composer.send());
+
+      await waitFor(() => {
+        expect(sent).toHaveBeenCalledWith(
+          expect.objectContaining({
+            threadId: aui.threads.getState().mainThreadId,
+          }),
+        );
+      });
+    } finally {
+      handle.destroy();
+    }
+  });
+
+  it("runs one chat per thread and keeps histories isolated across switches", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    expect(aui.threads.getState().threadIds).toEqual(["main"]);
+
+    flushTapSync(() => aui.composer.setText("first question"));
+    flushTapSync(() => aui.composer.send());
+    await waitFor(() => {
+      expect(
+        handle.getClient().thread.getState().messages.length,
+      ).toBeGreaterThan(0);
+    });
+    emit(...textReply("first answer"));
+    close();
+    await waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "first question",
+        "first answer",
+      ]);
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+    });
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    const state = handle.getClient().threads.getState();
+    expect(state.threadIds).toHaveLength(2);
+    expect(state.mainThreadId).not.toBe("main");
+    expect(handle.getClient().thread.getState().messages).toHaveLength(0);
+
+    flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+    expect(handle.getClient().threads.getState().mainThreadId).toBe("main");
+    expect(threadText(handle as never)).toEqual([
+      "first question",
+      "first answer",
+    ]);
+
+    handle.destroy();
+  });
+
+  it("keeps a switched-away thread streaming in the background", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport: () => transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    flushTapSync(() => aui.composer.setText("stream me"));
+    flushTapSync(() => aui.composer.send());
+    await waitFor(() => {
+      expect(
+        handle.getClient().thread.getState().messages.length,
+      ).toBeGreaterThan(0);
+    });
+    emit(
+      { type: "start" },
+      { type: "text-start", id: "t1" },
+      { type: "text-delta", id: "t1", delta: "partial" },
+    );
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    expect(handle.getClient().thread.getState().messages).toHaveLength(0);
+
+    emit(
+      { type: "text-delta", id: "t1", delta: " and finished" },
+      { type: "text-end", id: "t1" },
+      { type: "finish" },
+    );
+    close();
+
+    flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+    await waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "stream me",
+        "partial and finished",
+      ]);
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+    });
+
+    handle.destroy();
+  });
+
+  describe("a frontend tool running across a thread switch", () => {
+    const switchAwayAndBack = (
+      handle: ReturnType<typeof createAssistantClient>,
+    ) => {
+      flushTapSync(() => handle.getClient().threads.switchToNewThread());
+      flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+    };
+
+    it("executes a tool call streamed while switched away once after switching back", async () => {
+      const execute = vi.fn(() => new Promise<string>(() => {}));
+      const { transport, emit, close } = createControlledTransport();
+      const sendMessages = vi.spyOn(transport, "sendMessages");
+      const onToolCall = vi.fn();
+      const handle = createAssistantClient(
+        AuiConfig({ threads: AISDKThreads({ transport, onToolCall }) }),
+      );
+      handle.subscribe(() => {});
+      try {
+        flushTapSync(() =>
+          handle.getClient().modelContext.register({
+            getModelContext: () => ({
+              tools: {
+                lookup: {
+                  parameters: { type: "object", properties: {} },
+                  execute,
+                },
+              },
+            }),
+          }),
+        );
+        flushTapSync(() => handle.getClient().composer.setText("look up"));
+        flushTapSync(() => handle.getClient().composer.send());
+        await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(1));
+
+        emit({ type: "start", messageId: "assistant-1" });
+        await waitFor(() =>
+          expect(
+            handle
+              .getClient()
+              .thread.getState()
+              .messages.some((message) => message.role === "assistant"),
+          ).toBe(true),
+        );
+        flushTapSync(() => handle.getClient().threads.switchToNewThread());
+        emit(
+          {
+            type: "tool-input-available",
+            toolCallId: "call-after-switch",
+            toolName: "lookup",
+            input: {},
+          },
+          { type: "finish" },
+        );
+        await waitFor(() => expect(onToolCall).toHaveBeenCalledOnce());
+        expect(execute).not.toHaveBeenCalled();
+
+        flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+        close();
+        await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+        expect(handle.getClient().thread.getState().isRunning).toBe(true);
+      } finally {
+        handle.destroy();
+      }
+    });
+
+    it("keeps the thread running after switching back, and Stop aborts it", async () => {
+      let signal: AbortSignal | undefined;
+      const execute = vi.fn(
+        (_args: unknown, { abortSignal }: { abortSignal: AbortSignal }) =>
+          new Promise<string>(() => {
+            signal = abortSignal;
+          }),
+      );
+      const { handle } = await startTool(execute);
+      try {
+        await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+
+        switchAwayAndBack(handle);
+        expect(handle.getClient().thread.getState().isRunning).toBe(true);
+        flushTapSync(() => handle.getClient().thread.cancelRun());
+        expect(signal?.aborted).toBe(true);
+      } finally {
+        handle.destroy();
+      }
+    });
+
+    it("still continues the turn with its result when nobody stops it", async () => {
+      let finishTool!: (value: string) => void;
+      const execute = vi.fn(
+        () =>
+          new Promise<string>((resolve) => {
+            finishTool = resolve;
+          }),
+      );
+      const { handle, sendMessages } = await startTool(execute);
+      try {
+        await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+
+        switchAwayAndBack(handle);
+        finishTool("done");
+        await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(2));
+      } finally {
+        handle.destroy();
+      }
+    });
+
+    it("continues the turn when it settles while the thread is in the background", async () => {
+      let finishTool!: (value: string) => void;
+      const execute = vi.fn(
+        () =>
+          new Promise<string>((resolve) => {
+            finishTool = resolve;
+          }),
+      );
+      const { handle, sendMessages } = await startTool(execute);
+      try {
+        await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+
+        flushTapSync(() => handle.getClient().threads.switchToNewThread());
+        finishTool("done");
+        await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(2));
+      } finally {
+        handle.destroy();
+      }
+    });
+
+    it("stays resumable when it waits on human input", async () => {
+      const execute = vi.fn(
+        async (
+          _args: unknown,
+          { human }: { human: (p: unknown) => unknown },
+        ) => ({ answer: await human({ question: "proceed?" }) }),
+      );
+      const { handle, sendMessages } = await startTool(execute);
+      try {
+        await waitFor(() => expect(execute).toHaveBeenCalledOnce());
+
+        switchAwayAndBack(handle);
+        expect(() =>
+          handle
+            .getClient()
+            .thread.message({ index: 1 })
+            .part({ toolCallId: "call-1" })
+            .resumeToolCall("yes"),
+        ).not.toThrow();
+        await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(2));
+      } finally {
+        handle.destroy();
+      }
+    });
+  });
+
+  it("stops all in-flight chats when its client is destroyed", async () => {
+    const chats: ReturnType<typeof createCancellableTransport>[] = [];
+    const handle = createAssistantClient(
+      AuiConfig({
+        threads: AISDKThreads({
+          transport: () => {
+            const chat = createCancellableTransport();
+            chats.push(chat);
+            return chat.transport;
+          },
+        }),
+      }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    try {
+      flushTapSync(() => aui.composer.setText("first"));
+      flushTapSync(() => aui.composer.send());
+      await waitFor(() => {
+        expect(aui.thread.getState().isRunning).toBe(true);
+      });
+
+      flushTapSync(() => aui.threads.switchToNewThread());
+      flushTapSync(() => handle.getClient().composer.setText("second"));
+      flushTapSync(() => handle.getClient().composer.send());
+      await waitFor(() => {
+        expect(chats).toHaveLength(2);
+        expect(handle.getClient().thread.getState().isRunning).toBe(true);
+      });
+    } finally {
+      handle.destroy();
+    }
+
+    await waitFor(() => {
+      expect(chats.map((chat) => chat.getCancelCount())).toEqual([1, 1]);
+    });
+  });
+
+  it("wires model context and per-thread transports onto the wire", async () => {
+    const bodies: unknown[] = [];
+    const fetchStub = vi.fn(async (_url: unknown, init?: RequestInit) => {
+      bodies.push(JSON.parse(String(init?.body)));
+      return new Response(
+        new ReadableStream({
+          start(c) {
+            c.enqueue(
+              new TextEncoder().encode(
+                'data: {"type":"start"}\n\ndata: [DONE]\n\n',
+              ),
+            );
+            c.close();
+          },
+        }),
+        { headers: { "Content-Type": "text/event-stream" } },
+      );
+    });
+    vi.stubGlobal("fetch", fetchStub);
+    try {
+      const transports: AssistantChatTransport<never>[] = [];
+      const handle = createAssistantClient(
+        AuiConfig({
+          threads: AISDKThreads({
+            transport: () => {
+              const transport = new AssistantChatTransport();
+              transports.push(transport as AssistantChatTransport<never>);
+              return transport;
+            },
+          }),
+        }),
+      );
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      flushTapSync(() =>
+        aui.modelContext.register({
+          getModelContext: () => ({ system: "wired system prompt" }),
+        }),
+      );
+      flushTapSync(() => aui.composer.setText("hello"));
+      flushTapSync(() => aui.composer.send());
+      await waitFor(() => expect(fetchStub).toHaveBeenCalledTimes(1));
+      expect(bodies[0]).toMatchObject({ system: "wired system prompt" });
+
+      flushTapSync(() => aui.threads.switchToNewThread());
+      await waitFor(() => expect(transports.length).toBe(2));
+      expect(transports[0]).not.toBe(transports[1]);
+
+      handle.destroy();
+    } finally {
+      vi.unstubAllGlobals();
+    }
+  });
+
+  it("bridges model context on the default transport path", async () => {
+    const bodies: unknown[] = [];
+    const fetchStub = vi.fn(async (_url: unknown, init?: RequestInit) => {
+      bodies.push(JSON.parse(String(init?.body)));
+      return new Response(
+        new ReadableStream({
+          start(c) {
+            c.enqueue(
+              new TextEncoder().encode(
+                'data: {"type":"start"}\n\ndata: [DONE]\n\n',
+              ),
+            );
+            c.close();
+          },
+        }),
+        { headers: { "Content-Type": "text/event-stream" } },
+      );
+    });
+    vi.stubGlobal("fetch", fetchStub);
+    try {
+      const handle = createAssistantClient(
+        AuiConfig({ threads: AISDKThreads() }),
+      );
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      flushTapSync(() =>
+        aui.modelContext.register({
+          getModelContext: () => ({ system: "default path system" }),
+        }),
+      );
+      flushTapSync(() => aui.composer.setText("hello"));
+      flushTapSync(() => aui.composer.send());
+      await waitFor(() => expect(fetchStub).toHaveBeenCalledTimes(1));
+      expect(bodies[0]).toMatchObject({ system: "default path system" });
+
+      handle.destroy();
+    } finally {
+      vi.unstubAllGlobals();
+    }
+  });
+
+  it("forwards ChatInit callbacks to each thread's chat from the latest render", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const sendMessages = vi.spyOn(transport, "sendMessages");
+    const mounted = chatCallbacks();
+    const latest = chatCallbacks();
+    latest.sendAutomaticallyWhen.mockReturnValueOnce(true);
+    let callbacks = mounted;
+    const { handle, rerender } = createLiveHandle(() => ({
+      transport: () => transport,
+      ...callbacks,
+    }));
+    const aui = handle.getClient();
+
+    callbacks = latest;
+    rerender();
+
+    flushTapSync(() => aui.composer.setText("hi"));
+    flushTapSync(() => aui.composer.send());
+    await waitFor(() => {
+      expect(
+        handle.getClient().thread.getState().messages.length,
+      ).toBeGreaterThan(0);
+    });
+    emit(
+      { type: "start" },
+      {
+        type: "tool-input-available",
+        toolCallId: "call-1",
+        toolName: "lookup",
+        input: {},
+      },
+      { type: "data-note", data: "note" },
+      { type: "finish" },
+    );
+    close();
+    await waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(2));
+    emit({ type: "start" }, { type: "error", errorText: "boom" });
+    close();
+    await waitFor(() => expect(latest.onFinish).toHaveBeenCalledTimes(2));
+
+    expect(latest.onToolCall).toHaveBeenCalledExactlyOnceWith({
+      toolCall: expect.objectContaining({ toolCallId: "call-1" }),
+    });
+    expect(latest.onData).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({ type: "data-note", data: "note" }),
+    );
+    expect(latest.onError).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({ message: "boom" }),
+    );
+    expect(latest.sendAutomaticallyWhen).toHaveBeenCalledOnce();
+    for (const callback of Object.values(mounted)) {
+      expect(callback).not.toHaveBeenCalled();
+    }
+
+    handle.destroy();
+  });
+
+  it("forwards the latest callbacks to a switched-away thread still streaming in the background", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const onFinishA = vi.fn();
+    const onFinishB = vi.fn();
+    let onFinish = onFinishA;
+    const { handle, rerender } = createLiveHandle(() => ({
+      transport: () => transport,
+      onFinish,
+    }));
+    const aui = handle.getClient();
+
+    flushTapSync(() => aui.composer.setText("stream me"));
+    flushTapSync(() => aui.composer.send());
+    await waitFor(() => {
+      expect(
+        handle.getClient().thread.getState().messages.length,
+      ).toBeGreaterThan(0);
+    });
+    emit(
+      { type: "start" },
+      { type: "text-start", id: "t1" },
+      { type: "text-delta", id: "t1", delta: "partial" },
+    );
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    onFinish = onFinishB;
+    rerender();
+
+    emit({ type: "text-end", id: "t1" }, { type: "finish" });
+    close();
+    await waitFor(() => expect(onFinishB).toHaveBeenCalledTimes(1));
+    expect(onFinishA).not.toHaveBeenCalled();
+
+    handle.destroy();
+  });
+
+  it("posts each thread's own id as the chat id", async () => {
+    const bodies: unknown[] = [];
+    const fetchStub = vi.fn(async (_url: unknown, init?: RequestInit) => {
+      bodies.push(JSON.parse(String(init?.body)));
+      return new Response(
+        new ReadableStream({
+          start(c) {
+            c.enqueue(
+              new TextEncoder().encode(
+                'data: {"type":"start"}\n\ndata: [DONE]\n\n',
+              ),
+            );
+            c.close();
+          },
+        }),
+        { headers: { "Content-Type": "text/event-stream" } },
+      );
+    });
+    vi.stubGlobal("fetch", fetchStub);
+    try {
+      const handle = createAssistantClient(
+        AuiConfig({ threads: AISDKThreads() }),
+      );
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      flushTapSync(() => aui.composer.setText("hello"));
+      flushTapSync(() => aui.composer.send());
+      await waitFor(() => expect(fetchStub).toHaveBeenCalledTimes(1));
+      expect(bodies[0]).toMatchObject({ id: "main" });
+
+      flushTapSync(() => aui.threads.switchToNewThread());
+      const newThreadId = handle.getClient().threads.getState().mainThreadId;
+      flushTapSync(() => handle.getClient().composer.setText("hi again"));
+      flushTapSync(() => handle.getClient().composer.send());
+      await waitFor(() => expect(fetchStub).toHaveBeenCalledTimes(2));
+      expect(bodies[1]).toMatchObject({ id: newThreadId });
+
+      handle.destroy();
+    } finally {
+      vi.unstubAllGlobals();
+    }
+  });
+});
+
+describe("AISDKThreads branch isolation", () => {
+  const completeRound = async (
+    handle: ReturnType<typeof createAssistantClient>,
+    emit: (...chunks: never[]) => void,
+    close: () => void,
+    question: string,
+    answer: string,
+  ) => {
+    const aui = handle.getClient();
+    flushTapSync(() => aui.composer.setText(question));
+    flushTapSync(() => aui.composer.send());
+    await waitFor(() => {
+      expect(aui.thread.getState().messages.length).toBeGreaterThan(0);
+    });
+    emit(...(textReply(answer) as never[]));
+    close();
+    await waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+    });
+  };
+
+  it("keeps switching between populated threads free of cross-thread branches", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "thread a question",
+      "thread a answer",
+    );
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "thread b question",
+      "thread b answer",
+    );
+
+    flushTapSync(() => aui.threads.switchToThread("main"));
+    await waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "thread a question",
+        "thread a answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 0 }).getState().branchCount).toBe(1);
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(1);
+
+    handle.destroy();
+  });
+
+  it("preserves intra-thread branches across a switch through an empty thread", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+
+    await completeRound(
+      handle,
+      emit as never,
+      close,
+      "branchy question",
+      "first answer",
+    );
+
+    flushTapSync(() => aui.thread.message({ index: 1 }).reload());
+    await waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(true);
+    });
+    emit(...(textReply("second answer") as never[]));
+    close();
+    await waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(false);
+      expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+    });
+
+    flushTapSync(() => aui.threads.switchToNewThread());
+    expect(handle.getClient().thread.getState().messages).toHaveLength(0);
+
+    flushTapSync(() => handle.getClient().threads.switchToThread("main"));
+    await waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "branchy question",
+        "second answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+
+    handle.destroy();
+  });
+});