@assistant-ui/ai-sdk 0.0.3 → 0.0.4

package/src/runtime/AISDKThreads.cloud.test.ts

@@ -95,7 +95,7 @@ describe("AISDKThreads cloud", () => {
     handle.destroy();
   });
 
-  it("stops an in-flight cloud chat when switching away", async () => {
+  it("keeps an in-flight cloud chat running across a switch and stops it on delete", async () => {
     const chat = createCancellableTransport();
     const handle = createAssistantClient(
       AuiConfig({
@@ -126,6 +126,16 @@ describe("AISDKThreads cloud", () => {
       await vi.waitFor(() => {
         expect(handle.getClient().threads.getState().mainThreadId).toBe("t2");
       });
+      expect(chat.getCancelCount()).toBe(0);
+      await vi.waitFor(() => {
+        expect(
+          handle.getClient().threads.item({ id: "t1" }).getState().isRunning,
+        ).toBe(true);
+      });
+
+      flushTapSync(() =>
+        handle.getClient().threads.item({ id: "t1" }).delete(),
+      );
       await vi.waitFor(() => {
         expect(chat.getCancelCount()).toBe(1);
       });

package/src/runtime/AISDKThreads.test.ts

@@ -546,3 +546,105 @@ describe("AISDKThreads", () => {
     }
   });
 });
+
+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 vi.waitFor(() => {
+      expect(aui.thread.getState().messages.length).toBeGreaterThan(0);
+    });
+    emit(...(textReply(answer) as never[]));
+    close();
+    await vi.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 vi.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 vi.waitFor(() => {
+      expect(handle.getClient().thread.getState().isRunning).toBe(true);
+    });
+    emit(...(textReply("second answer") as never[]));
+    close();
+    await vi.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 vi.waitFor(() => {
+      expect(threadText(handle as never)).toEqual([
+        "branchy question",
+        "second answer",
+      ]);
+    });
+    expect(aui.thread.message({ index: 1 }).getState().branchCount).toBe(2);
+
+    handle.destroy();
+  });
+});

package/src/runtime/AISDKThreads.ts

@@ -24,6 +24,7 @@ import {
   useChatThread,
   type ChatThreadOptions,
 } from "./useChatThread";
+import { MessageRepository } from "@assistant-ui/core/internal";
 import { useResourceCleanup } from "./useResourceCleanup";
 
 export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
@@ -41,9 +42,9 @@ export type AISDKThreadsOptions<UI_MESSAGE extends UIMessage = UIMessage> =
       | undefined;
     /**
      * When set, the thread list is a `RemoteThreadList` backed by this
-     * assistant-cloud. Omit it to keep the in-memory list. The thread
-     * factory is keyed so cloud history reloads on a switch, and an
-     * in-flight run does not continue in the background.
+     * assistant-cloud. Omit it to keep the in-memory list. Every visited
+     * cloud thread stays mounted, so an in-flight run continues after a
+     * switch and stops on delete; per-thread history loads once per thread.
      */
     cloud?: AssistantCloud | undefined;
     /**
@@ -64,6 +65,7 @@ type AISDKThreadChatOptions<UI_MESSAGE extends UIMessage = UIMessage> = Omit<
 type ChatEntry<UI_MESSAGE extends UIMessage> = {
   chat: Chat<UI_MESSAGE>;
   transport: ChatTransport<UI_MESSAGE>;
+  repository: MessageRepository;
 };
 
 const createChatEntry = <UI_MESSAGE extends UIMessage>(
@@ -84,6 +86,7 @@ const createChatEntry = <UI_MESSAGE extends UIMessage>(
   return {
     chat: new Chat<UI_MESSAGE>({ ...chatInit, id: threadId, transport }),
     transport,
+    repository: new MessageRepository(),
   };
 };
 
@@ -113,7 +116,7 @@ const useAISDKChatThread = <UI_MESSAGE extends UIMessage = UIMessage>({
   const [owned] = useState(() =>
     cloud ? createChatEntry(threadId, options) : undefined,
   );
-  const { chat, transport } =
+  const { chat, transport, repository } =
     owned ?? getOrCreateChatEntry(threadId, options, chats);
 
   useEffect(() => {
@@ -145,6 +148,7 @@ const useAISDKChatThread = <UI_MESSAGE extends UIMessage = UIMessage>({
             : undefined
           : fallbackItem,
       chat,
+      messageRepositoryInstance: repository,
       stopOnClientDestroy: cloud,
     },
   );
@@ -190,6 +194,7 @@ const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
     bindCloud
       ? RemoteThreadList({
           adapter: cloudAdapter,
+          backgroundThreads: true,
           thread,
           threadId,
           onThreadIdChange,
@@ -212,11 +217,13 @@ const useAISDKThreads = <UI_MESSAGE extends UIMessage = UIMessage>(
  * per-thread orchestration as {@link AISDKChat} inside the client's own
  * resource tree, so it works with any `AssistantClient` host, React or not.
  * Without `cloud`, threads live in memory for the client's lifetime and keep
- * their history across switches; each thread's chat id is its thread id.
- * With `cloud`, the list is a `RemoteThreadList` and the factory is keyed so
- * cloud history reloads on a switch. The store entry mounts only the visible
- * thread, so a switch cancels an in-flight run. Model context is
- * registered on the visible thread only.
+ * their history across switches; only the visible thread is mounted, and a
+ * switched-away chat keeps streaming into its stored state until it settles
+ * or the thread is deleted. With `cloud`, the list is a
+ * `RemoteThreadList` with `backgroundThreads`: every visited thread stays
+ * mounted with its own history, a run continues after a switch and stops on
+ * delete, and a freshly created thread titles itself. Model context is
+ * registered on every mounted thread.
  */
 export const AISDKThreads = resource(useAISDKThreads);
 

package/src/runtime/useAISDKRuntime.test.ts

@@ -27,6 +27,7 @@ const createChatHelpers = (messages: any[] = []) => {
   let currentMessages = [...messages];
 
   const chatHelpers: any = {
+    id: "chat-1",
     status: "ready",
     error: null,
     messages: currentMessages,
@@ -135,10 +136,11 @@ describe("useAISDKRuntime", () => {
     abortError.name = "AbortError";
     const chat = createChatHelpers();
     let stopCalls = 0;
+    let rejectStop!: (error: unknown) => void;
     chat.stop = () => {
       stopCalls += 1;
       return new Promise((_, reject) => {
-        setTimeout(() => reject(abortError), 5);
+        rejectStop = reject;
       });
     };
     const consoleError = vi
@@ -147,8 +149,11 @@ describe("useAISDKRuntime", () => {
 
     try {
       const { result } = renderHook(() => useAISDKRuntime(chat));
-      const unhandledRejections = await captureUnhandledRejections(() => {
-        result.current.thread.cancelRun();
+      const unhandledRejections = await captureUnhandledRejections(async () => {
+        await act(async () => {
+          result.current.thread.cancelRun();
+          rejectStop(abortError);
+        });
       });
 
       expect(stopCalls).toBe(1);
@@ -159,6 +164,286 @@ describe("useAISDKRuntime", () => {
     }
   });
 
+  it("marks only the stopped output cancelled", async () => {
+    let resolveStop!: () => void;
+    const chat = createChatHelpers([
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "partial", state: "streaming" }],
+      },
+    ]);
+    chat.status = "streaming";
+    chat.stop = vi.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveStop = resolve;
+        }),
+    );
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    act(() => {
+      result.current.thread.cancelRun();
+    });
+    rerender();
+
+    expect(
+      result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({ type: "running" });
+
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({
+        type: "incomplete",
+        reason: "cancelled",
+      });
+    });
+
+    await act(async () => {
+      resolveStop();
+      await Promise.resolve();
+    });
+
+    act(() => {
+      chat.setMessages([
+        {
+          id: "assistant-2",
+          role: "assistant",
+          parts: [{ type: "text", text: "replacement" }],
+        },
+      ]);
+      rerender();
+    });
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+
+    act(() => {
+      chat.setMessages([
+        {
+          id: "assistant-1",
+          role: "assistant",
+          parts: [{ type: "text", text: "partial" }],
+        },
+      ]);
+      chat.status = "streaming";
+      rerender();
+    });
+
+    expect(
+      result.current.thread.getState().messages.at(-1)?.status,
+    ).toMatchObject({ type: "running" });
+
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({
+        type: "complete",
+        reason: "unknown",
+      });
+    });
+  });
+
+  it("keeps the stopped output cancelled through the next turn", async () => {
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "hi" }] },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "partial", state: "streaming" }],
+      },
+    ]);
+    chat.status = "streaming";
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    await act(async () => {
+      result.current.thread.cancelRun();
+    });
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+    });
+
+    act(() => {
+      chat.setMessages([
+        ...chat.messages,
+        { id: "u2", role: "user", parts: [{ type: "text", text: "next" }] },
+        {
+          id: "assistant-2",
+          role: "assistant",
+          parts: [{ type: "text", text: "answer" }],
+        },
+      ]);
+      rerender();
+    });
+
+    await waitFor(() => {
+      const messages = result.current.thread.getState().messages;
+      expect(
+        messages.find((message) => message.id === "assistant-1")?.status,
+      ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+      expect(messages.at(-1)?.status).toMatchObject({
+        type: "complete",
+        reason: "unknown",
+      });
+    });
+  });
+
+  it("retracts the cancellation when the provider picks the message back up", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "partial", state: "streaming" }],
+      },
+    ]);
+    chat.status = "streaming";
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    await act(async () => {
+      result.current.thread.cancelRun();
+    });
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+    });
+
+    act(() => {
+      chat.status = "streaming";
+      rerender();
+    });
+    act(() => {
+      chat.status = "ready";
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+  });
+
+  it("does not mark completed output cancelled when already idle", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "finished" }],
+      },
+    ]);
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+
+    act(() => {
+      result.current.thread.cancelRun();
+      rerender();
+    });
+
+    await waitFor(() => {
+      expect(
+        result.current.thread.getState().messages.at(-1)?.status,
+      ).toMatchObject({ type: "complete", reason: "unknown" });
+    });
+  });
+
+  it("marks output cancelled while a client tool is still executing", async () => {
+    let resolveTool!: (value: string) => void;
+    const execute = vi.fn(
+      () =>
+        new Promise<string>((resolve) => {
+          resolveTool = resolve;
+        }),
+    );
+    const chat = createChatHelpers();
+
+    const { result, rerender } = renderHook(() => useAISDKRuntime(chat));
+    const unregister = result.current.registerModelContextProvider({
+      getModelContext: () => ({
+        tools: {
+          weather: {
+            parameters: { type: "object", properties: {} },
+            execute,
+          },
+        },
+      }),
+    });
+
+    try {
+      act(() => {
+        chat.setMessages([
+          {
+            id: "assistant-1",
+            role: "assistant",
+            parts: [
+              {
+                type: "tool-weather",
+                toolCallId: "tool-1",
+                state: "input-available",
+                input: { city: "London" },
+              },
+            ],
+          },
+        ]);
+        rerender();
+      });
+
+      await waitFor(() => {
+        expect(execute).toHaveBeenCalledOnce();
+        expect(result.current.thread.getState().isRunning).toBe(true);
+      });
+
+      act(() => {
+        result.current.thread.cancelRun();
+        chat.setMessages([
+          {
+            id: "assistant-1",
+            role: "assistant",
+            parts: [{ type: "text", text: "stopped" }],
+          },
+        ]);
+        resolveTool("sunny");
+        rerender();
+      });
+
+      await waitFor(() => {
+        expect(
+          result.current.thread.getState().messages.at(-1)?.status,
+        ).toMatchObject({ type: "incomplete", reason: "cancelled" });
+      });
+    } finally {
+      unregister();
+    }
+  });
+
   it("reports non-AbortError cancellation failures", async () => {
     const stopError = new Error("stop failed");
     const chat = createChatHelpers();

package/src/runtime/useAISDKRuntime.ts

@@ -1,6 +1,13 @@
 "use client";
 
-import { useEffect, useMemo, useRef, useState } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
 import type {
   UIMessage,
   useChat,
@@ -82,6 +89,7 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
         })
       | undefined;
     toCreateMessage?: CustomToCreateMessageFunction;
+    unstable_messageRepositoryInstance?: MessageRepository | undefined;
     /**
      * Whether to automatically cancel pending interactive tool calls when the user sends a new message.
      *
@@ -142,9 +150,13 @@ const useGeneratedSuggestions = (
   const controllerRef = useRef<AbortController | null>(null);
   const wasRunningRef = useRef(false);
   const messagesRef = useRef(messages);
-  messagesRef.current = messages;
+  useInsertionEffect(() => {
+    messagesRef.current = messages;
+  }, [messages]);
   const adapterRef = useRef(suggestionAdapter);
-  adapterRef.current = suggestionAdapter;
+  useInsertionEffect(() => {
+    adapterRef.current = suggestionAdapter;
+  }, [suggestionAdapter]);
   const hasAdapter = suggestionAdapter != null;
 
   useEffect(() => {
@@ -205,6 +217,8 @@ const useGeneratedSuggestions = (
   return suggestions;
 };
 
+const NO_CANCELLED_MESSAGE_IDS: ReadonlySet<string> = new Set();
+
 export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   chatHelpers: ReturnType<typeof useChat<UI_MESSAGE>>,
   adapter: AISDKRuntimeAdapter<UI_MESSAGE> = {},
@@ -224,6 +238,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const [toolStatuses, setToolStatuses] = useState<
     Record<string, ToolExecutionStatus>
   >({});
+  const [cancelledMessages, setCancelledMessages] = useState<{
+    chatId: string;
+    ids: ReadonlySet<string>;
+  } | null>(null);
   const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
     new Map(),
   );
@@ -239,6 +257,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const providerIsRunning =
     chatHelpers.status === "submitted" || chatHelpers.status === "streaming";
   const isRunning = providerIsRunning || hasExecutingTools;
+  const wasProviderRunningRef = useRef(providerIsRunning);
 
   const messageTiming = useStreamingTiming(chatHelpers.messages, isRunning);
 
@@ -248,6 +267,42 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const optimisticMessageId =
     isRunning && lastMessage?.role === "assistant" ? lastMessage.id : undefined;
 
+  const cancelledMessageIds =
+    cancelledMessages?.chatId === chatHelpers.id
+      ? cancelledMessages.ids
+      : NO_CANCELLED_MESSAGE_IDS;
+
+  const retractCancellation = useCallback(
+    (chatId: string, messageId: string) => {
+      setCancelledMessages((prev) => {
+        if (prev?.chatId !== chatId || !prev.ids.has(messageId)) return prev;
+        const ids = new Set(prev.ids);
+        ids.delete(messageId);
+        return { chatId, ids };
+      });
+    },
+    [],
+  );
+
+  // A provider run that resumes the stopped response retracts its cancellation;
+  // a run that starts a new response leaves the stopped one marked.
+  const resumedMessageId =
+    providerIsRunning && lastMessage?.role === "assistant"
+      ? lastMessage.id
+      : undefined;
+
+  useEffect(() => {
+    const wasProviderRunning = wasProviderRunningRef.current;
+    wasProviderRunningRef.current = providerIsRunning;
+    if (wasProviderRunning || !resumedMessageId) return;
+    retractCancellation(chatHelpers.id, resumedMessageId);
+  }, [
+    providerIsRunning,
+    resumedMessageId,
+    chatHelpers.id,
+    retractCancellation,
+  ]);
+
   const messages = AISDKMessageConverter.useThreadMessages({
     isRunning,
     messages: chatHelpers.messages,
@@ -261,8 +316,15 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         mcpAppMetadataCache: mcpAppMetadataCacheRef.current,
         ...(optimisticMessageId && { optimisticMessageId }),
         ...(chatHelpers.error && { error: chatHelpers.error.message }),
+        ...(cancelledMessageIds.size > 0 && { cancelledMessageIds }),
       }),
-      [toolStatuses, messageTiming, optimisticMessageId, chatHelpers.error],
+      [
+        toolStatuses,
+        messageTiming,
+        optimisticMessageId,
+        chatHelpers.error,
+        cancelledMessageIds,
+      ],
     ),
   });
 
@@ -432,10 +494,27 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
       runtimeRef.current.thread.import(exportedRepo);
     },
     onCancel: async () => {
+      const message = chatHelpers.messages.at(-1);
+      const cancelledId =
+        isRunning && message?.role === "assistant" ? message.id : undefined;
+      if (cancelledId) {
+        const liveIds = new Set(chatHelpers.messages.map((m) => m.id));
+        setCancelledMessages((prev) => {
+          const kept =
+            prev?.chatId === chatHelpers.id
+              ? [...prev.ids].filter((id) => liveIds.has(id))
+              : [];
+          return {
+            chatId: chatHelpers.id,
+            ids: new Set([...kept, cancelledId]),
+          };
+        });
+      }
       try {
         await chatHelpers.stop();
       } catch (error) {
         if (!(error instanceof Error && error.name === "AbortError")) {
+          if (cancelledId) retractCancellation(chatHelpers.id, cancelledId);
           throw error;
         }
       }
@@ -549,6 +628,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         }),
       ),
     ...pickExternalStoreSharedOptions(adapter),
+    ...(adapter.unstable_messageRepositoryInstance && {
+      unstable_messageRepositoryInstance:
+        adapter.unstable_messageRepositoryInstance,
+    }),
     ...(suggestionAdapter ? { suggestions: generatedSuggestions } : {}),
     ...(onResume && { onResume }),
     ...(onResumeToolCall && { onResumeToolCall }),
@@ -563,7 +646,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   });
 
   const setMessagesRef = useRef(chatHelpers.setMessages);
-  setMessagesRef.current = chatHelpers.setMessages;
+  useInsertionEffect(() => {
+    setMessagesRef.current = chatHelpers.setMessages;
+  }, [chatHelpers.setMessages]);
 
   useEffect(() => {
     if (hasSeededRepositoryRef.current) return;