@assistant-ui/ai-sdk 0.0.8 → 0.0.10

package/src/runtime/useChatThread.ts

@@ -181,6 +181,8 @@ type ChatCallbacks<UI_MESSAGE extends UIMessage> = Pick<
   "onToolCall" | "onData" | "onFinish" | "onError" | "sendAutomaticallyWhen"
 >;
 
+const requestsByChat = new WeakMap<object, symbol>();
+
 /**
  * Constructs a `Chat` whose callbacks read the latest options through
  * `callbacksRef`, the forwarding `useChat` applies only to a chat it
@@ -189,9 +191,22 @@ type ChatCallbacks<UI_MESSAGE extends UIMessage> = Pick<
 export const createChat = <UI_MESSAGE extends UIMessage>(
   init: ChatInit<UI_MESSAGE>,
   callbacksRef: { readonly current: ChatCallbacks<UI_MESSAGE> | undefined },
-): Chat<UI_MESSAGE> =>
-  new Chat<UI_MESSAGE>({
+): Chat<UI_MESSAGE> => {
+  const transport = init.transport;
+  const chat = new Chat<UI_MESSAGE>({
     ...init,
+    ...(transport && {
+      transport: {
+        sendMessages: (options) => {
+          requestsByChat.set(chat, Symbol());
+          return transport.sendMessages(options);
+        },
+        reconnectToStream: (options) => {
+          requestsByChat.set(chat, Symbol());
+          return transport.reconnectToStream(options);
+        },
+      },
+    }),
     onToolCall: (arg) => callbacksRef.current?.onToolCall?.(arg),
     onData: (arg) => callbacksRef.current?.onData?.(arg),
     onFinish: (arg) => callbacksRef.current?.onFinish?.(arg),
@@ -199,6 +214,8 @@ export const createChat = <UI_MESSAGE extends UIMessage>(
     sendAutomaticallyWhen: (arg) =>
       callbacksRef.current?.sendAutomaticallyWhen?.(arg) ?? false,
   });
+  return chat;
+};
 
 export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   options: ChatThreadOptions<UI_MESSAGE> | undefined,
@@ -271,7 +288,10 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
   );
 
   const runtime = useAISDKRuntime(chat, {
-    adapters,
+    adapters: {
+      ...adapters,
+      threadList: { threadId: id, ...adapters?.threadList },
+    },
     ...pickExternalStoreSharedOptions(options ?? {}),
     ...(toCreateMessage && { toCreateMessage }),
     ...(onResume && { onResume }),
@@ -282,6 +302,11 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     ...(messageRepositoryInstance && {
       unstable_messageRepositoryInstance: messageRepositoryInstance,
     }),
+    // The chat outlives this runtime when a host mounts only the visible
+    // thread, so a host approval answer is kept with it. This is the Chat
+    // instance, not the useChat helpers, which are re-minted every render and
+    // would be a dead WeakMap key by the next one.
+    unstable_hostApprovalOwner: externalChat ?? ownedChat,
     ...(unstable_onBranchChange && { unstable_onBranchChange }),
   });
 
@@ -350,27 +375,41 @@ export const useChatThread = <UI_MESSAGE extends UIMessage = UIMessage>(
     }
     if (isLoadingHistory) return;
     resumedStreamIds.add(pendingStreamId);
-    chat.resumeStream().catch((err: unknown) => {
-      console.warn("[assistant-ui] resumable: resume failed", err);
-      try {
-        onResumeErrorRef.current?.(err);
-      } catch (callbackError) {
-        console.error(
-          "[assistant-ui] resumable: onResumeError callback failed",
-          callbackError,
-        );
-      } finally {
-        if (resumableStorage?.getStreamId(id) === pendingStreamId) {
-          resumableStorage.clear(id);
+    const activeChat = externalChat ?? ownedChat;
+    activeChat.clearError();
+    const pending = chat.resumeStream();
+    const request = requestsByChat.get(activeChat);
+    pending
+      .then(() => {
+        // Chat.error is shared with sends and resumes that can start before
+        // this promise settles, including inside the caller's onFinish.
+        if (requestsByChat.get(activeChat) === request && activeChat.error) {
+          throw activeChat.error;
+        }
+      })
+      .catch((err: unknown) => {
+        console.warn("[assistant-ui] resumable: resume failed", err);
+        try {
+          onResumeErrorRef.current?.(err);
+        } catch (callbackError) {
+          console.error(
+            "[assistant-ui] resumable: onResumeError callback failed",
+            callbackError,
+          );
+        } finally {
+          if (resumableStorage?.getStreamId(id) === pendingStreamId) {
+            resumableStorage.clear(id);
+          }
         }
-      }
-    });
+      });
   }, [
     chat,
+    externalChat,
     id,
     isChatRunning,
     isLoadingHistory,
     pendingStreamId,
+    ownedChat,
     resumableStorage,
     resumedStreamIds,
   ]);

package/src/runtime/useExternalHistory.ts

@@ -51,9 +51,20 @@ export const toExportedMessageRepository = <TMessage>(
   };
 };
 
+const hasUnansweredApproval = (message: ThreadMessage) =>
+  message.content.some(
+    (part) =>
+      part.type === "tool-call" &&
+      part.approval != null &&
+      part.approval.approved === undefined &&
+      part.approval.resolution === undefined,
+  );
+
+// Core reports a run paused on an unanswered approval as an interrupt, not as tool calls.
 const isAwaitingToolApproval = (message: ThreadMessage) =>
   message.status?.type === "requires-action" &&
-  message.status.reason === "tool-calls";
+  (message.status.reason === "tool-calls" ||
+    (message.status.reason === "interrupt" && hasUnansweredApproval(message)));
 
 const isTerminalMessage = (message: ThreadMessage) =>
   message.status === undefined ||

package/src/transport/AssistantChatTransport.test.ts

@@ -205,6 +205,170 @@ const wrappedFetchOf = (
   ).fetch;
 
 describe("AssistantChatTransport resumable fetch wrapper", () => {
+  it("keeps a replacement checkpoint set while preparing an older reconnect", async () => {
+    const storage = createMemoryStorage("stream-old");
+    let finishPrepare!: () => void;
+    const prepareReconnectToStreamRequest = vi.fn(async () => {
+      await new Promise<void>((resolve) => {
+        finishPrepare = resolve;
+      });
+      return { headers: { "x-custom": "retained" } };
+    });
+    const fetch = vi
+      .fn<typeof globalThis.fetch>()
+      .mockResolvedValue(new Response(null, { status: 204 }));
+    const transport = new AssistantChatTransport({
+      fetch,
+      prepareReconnectToStreamRequest,
+      resumable: { storage, resumeApi: (id) => `/api/resume/${id}` },
+    });
+    const pending = transport.reconnectToStream({ chatId: "thread" });
+    await vi.waitFor(() =>
+      expect(prepareReconnectToStreamRequest).toHaveBeenCalledOnce(),
+    );
+    storage.setStreamId("stream-new");
+    finishPrepare();
+    await expect(pending).resolves.toBeNull();
+    expect(fetch.mock.calls[0]?.[0]).toBe("/api/resume/stream-old");
+    expect(
+      Array.from(new Headers(fetch.mock.calls[0]?.[1]?.headers).entries()),
+    ).toEqual([["x-custom", "retained"]]);
+    expect(storage.getStreamId()).toBe("stream-new");
+  });
+
+  it.each([
+    { status: 204, replaceCheckpoint: false },
+    { status: 204, replaceCheckpoint: true },
+    { status: 404, replaceCheckpoint: false },
+    { status: 404, replaceCheckpoint: true },
+  ])(
+    "clears only the matching checkpoint after $status (replacement: $replaceCheckpoint)",
+    async ({ status, replaceCheckpoint }) => {
+      const storage = createMemoryStorage("stream-old");
+      let respond!: (response: Response) => void;
+      const fetch = vi.fn(
+        () =>
+          new Promise<Response>((resolve) => {
+            respond = resolve;
+          }),
+      );
+      const transport = new AssistantChatTransport({
+        fetch,
+        resumable: { storage, resumeApi: "/api/resume" },
+      });
+      const pending = transport.reconnectToStream({ chatId: "thread" });
+      await vi.waitFor(() => expect(fetch).toHaveBeenCalledOnce());
+      if (replaceCheckpoint) storage.setStreamId("stream-new");
+      respond(
+        new Response(status === 404 ? "stream expired" : null, { status }),
+      );
+      if (status === 404) {
+        await expect(pending).rejects.toThrow("stream expired");
+      } else {
+        await expect(pending).resolves.toBeNull();
+      }
+      expect(storage.getStreamId()).toBe(
+        replaceCheckpoint ? "stream-new" : null,
+      );
+    },
+  );
+
+  it.each([undefined, "stream-old", "response-id"])(
+    "preserves a replacement checkpoint on a delayed successful reconnect (%s)",
+    async (responseId) => {
+      const storage = createMemoryStorage("stream-old");
+      let respond!: (response: Response) => void;
+      const fetch = vi.fn(
+        () =>
+          new Promise<Response>((resolve) => {
+            respond = resolve;
+          }),
+      );
+      const transport = new AssistantChatTransport({
+        fetch,
+        resumable: { storage, resumeApi: "/api/resume" },
+      });
+      const pending = transport.reconnectToStream({ chatId: "thread" });
+      await vi.waitFor(() => expect(fetch).toHaveBeenCalledOnce());
+      storage.setStreamId("stream-new");
+      respond(
+        new Response('data: {"type":"finish"}\n\n', {
+          headers: {
+            "content-type": "text/event-stream",
+            ...(responseId && { [RESUMABLE_STREAM_ID_HEADER]: responseId }),
+          },
+        }),
+      );
+      const stream = await pending;
+      expect(storage.getStreamId()).toBe("stream-new");
+      const reader = stream!.getReader();
+      while (!(await reader.read()).done) {}
+      expect(storage.getStreamId()).toBe("stream-new");
+    },
+  );
+
+  it.each([undefined, "stream-old", "response-id"])(
+    "preserves a checkpoint replaced while consuming a reconnect (%s)",
+    async (responseId) => {
+      const storage = createMemoryStorage("stream-old");
+      let controller!: ReadableStreamDefaultController<Uint8Array>;
+      const transport = new AssistantChatTransport({
+        fetch: vi.fn(
+          async () =>
+            new Response(
+              new ReadableStream({
+                start(value) {
+                  controller = value;
+                },
+              }),
+              {
+                headers: {
+                  "content-type": "text/event-stream",
+                  ...(responseId && {
+                    [RESUMABLE_STREAM_ID_HEADER]: responseId,
+                  }),
+                },
+              },
+            ),
+        ),
+        resumable: { storage, resumeApi: "/api/resume" },
+      });
+      const stream = await transport.reconnectToStream({ chatId: "thread" });
+      expect(storage.getStreamId()).toBe(responseId ?? "stream-old");
+      storage.setStreamId("stream-new");
+      controller.enqueue(
+        new TextEncoder().encode('data: {"type":"finish"}\n\n'),
+      );
+      controller.close();
+      const reader = stream!.getReader();
+      while (!(await reader.read()).done) {}
+      expect(storage.getStreamId()).toBe("stream-new");
+    },
+  );
+
+  it.each([undefined, "stream-old", "response-id"])(
+    "clears the checkpoint owned by a completed reconnect (%s)",
+    async (responseId) => {
+      const storage = createMemoryStorage("stream-old");
+      const transport = new AssistantChatTransport({
+        fetch: vi.fn(
+          async () =>
+            new Response('data: {"type":"finish"}\n\n', {
+              headers: {
+                "content-type": "text/event-stream",
+                ...(responseId && { [RESUMABLE_STREAM_ID_HEADER]: responseId }),
+              },
+            }),
+        ),
+        resumable: { storage, resumeApi: "/api/resume" },
+      });
+      const stream = await transport.reconnectToStream({ chatId: "thread" });
+      const reader = stream!.getReader();
+      while (!(await reader.read()).done) {}
+      expect(storage.getStreamId()).toBeNull();
+    },
+  );
+
   it("passes a 204 with a non-null empty body through untouched (WebKit)", async () => {
     const response = nullBodyStatusWithBody(204);
     const fetchMock = vi.fn(async () => response);

package/src/transport/AssistantChatTransport.ts

@@ -22,6 +22,7 @@ const FINISH_MARKER = '"type":"finish"';
 const FINISH_BUFFER_LIMIT = 4096;
 const FINISH_BUFFER_TAIL = 1024;
 const RESUMABLE_THREAD_ID_HEADER = "x-assistant-ui-resumable-thread-id";
+const RESUMABLE_RECONNECT_ID_HEADER = "x-assistant-ui-resumable-reconnect-id";
 
 // 101/204/205/304 are null-body statuses per the fetch spec: `new Response(body, { status })`
 // throws for them, and WebKit returns a non-null empty body, so the body check alone does not guard it.
@@ -137,10 +138,24 @@ function wrapFetchWithResumable(
   return async (input, init) => {
     const headers = new Headers(init?.headers);
     const threadId = headers.get(RESUMABLE_THREAD_ID_HEADER) ?? undefined;
+    const reconnectingStreamId = headers.get(RESUMABLE_RECONNECT_ID_HEADER);
+    const checkpointAtRequest =
+      reconnectingStreamId ?? resumable.storage.getStreamId(threadId);
     headers.delete(RESUMABLE_THREAD_ID_HEADER);
+    headers.delete(RESUMABLE_RECONNECT_ID_HEADER);
     const res = await baseFetch(input, { ...init, headers });
     const id = res.headers.get(RESUMABLE_STREAM_ID_HEADER);
-    if (id) resumable.storage.setStreamId(id, threadId);
+    const ownsCheckpoint =
+      !reconnectingStreamId ||
+      resumable.storage.getStreamId(threadId) === reconnectingStreamId;
+    if (id && ownsCheckpoint) resumable.storage.setStreamId(id, threadId);
+    if (
+      (res.status === 204 || res.status === 404) &&
+      reconnectingStreamId &&
+      resumable.storage.getStreamId(threadId) === reconnectingStreamId
+    ) {
+      resumable.storage.clear(threadId);
+    }
     if (!res.body || NULL_BODY_STATUSES.has(res.status)) return res;
 
     const detectFinish = resumable.isFinishEvent ?? defaultIsFinishEvent;
@@ -153,7 +168,11 @@ function wrapFetchWithResumable(
         controller.enqueue(chunk);
         accumulator += decoder.decode(chunk, { stream: true });
         if (detectFinish(chunk, accumulator)) {
-          if (!id || resumable.storage.getStreamId(threadId) === id) {
+          if (
+            ownsCheckpoint &&
+            resumable.storage.getStreamId(threadId) ===
+              (id ?? checkpointAtRequest)
+          ) {
             resumable.storage.clear(threadId);
           }
           accumulator = "";
@@ -195,6 +214,7 @@ function wrapPrepareReconnect(
     const userPrepared = await userPrepareReconnect?.({ ...options, api });
     const headers = new Headers(userPrepared?.headers ?? options.headers);
     headers.set(RESUMABLE_THREAD_ID_HEADER, options.id);
+    headers.set(RESUMABLE_RECONNECT_ID_HEADER, streamId);
     return {
       ...userPrepared,
       headers,