@assistant-ui/core 0.3.20 → 0.3.21

package/src/react/runtimes/assistant-transport/replayBoundaryStream.test.ts

@@ -322,6 +322,7 @@ describe("createReplayBoundaryStream", () => {
 
   it("releases the response body reader after a read failure", async () => {
     const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
     const error = new Error("stream failed");
     const body = new ReadableStream<Uint8Array>({
       start(controller) {
@@ -330,13 +331,14 @@ describe("createReplayBoundaryStream", () => {
     });
     const streamPromise = createReplayBoundaryStream(
       new Response(body, { headers: { [REPLAY_CONTENT_LENGTH_HEADER]: "10" } }),
-      { setReplaying: vi.fn(), waitForRender },
+      { setReplaying, waitForRender },
     );
 
     await releaseNext();
     const stream = await streamPromise;
 
     await expect(stream.getReader().read()).rejects.toBe(error);
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
     expect(body.locked).toBe(false);
   });
 

package/src/react/runtimes/assistant-transport/replayBoundaryStream.ts

@@ -90,8 +90,15 @@ export const createReplayBoundaryStream = async (
   const reader = body.getReader();
   let bytesForwarded = 0;
   let replayFinished = false;
+  let replayCleared = false;
   let readerCleanup: Promise<void> | undefined;
 
+  const clearReplay = () => {
+    if (replayCleared) return;
+    replayCleared = true;
+    setReplaying(false);
+  };
+
   const releaseReader = () => {
     if (readerCleanup) return readerCleanup;
     reader.releaseLock();
@@ -117,7 +124,7 @@ export const createReplayBoundaryStream = async (
 
     // Let replay bytes drain before rendering live mode, then render live mode before releasing live bytes.
     await waitForReplayRender();
-    setReplaying(false);
+    clearReplay();
     await waitForReplayRender();
   };
 
@@ -158,6 +165,7 @@ export const createReplayBoundaryStream = async (
         await finishReplay();
         controller.enqueue(value.subarray(replayBytesInChunk));
       } catch (error) {
+        clearReplay();
         await cancelReader(error).catch(() => {});
         throw error;
       }
@@ -165,7 +173,7 @@ export const createReplayBoundaryStream = async (
     async cancel(reason) {
       const wasFinished = replayFinished;
       replayFinished = true;
-      if (!wasFinished) setReplaying(false);
+      if (!wasFinished) clearReplay();
       await cancelReader(reason);
     },
   });

package/src/react/runtimes/assistant-transport/runManager.ts

@@ -49,6 +49,9 @@ export function useRunManager(config: {
       try {
         if (!disposeAborted()) {
           await onRunRef.current(ac.signal);
+          // A fully received body is not errored by abort(), so a cancelled
+          // run can still resolve.
+          if (ac.signal.aborted) throw ac.signal.reason;
         }
       } catch (error) {
         if (!disposeAborted() && !stateRef.current.disposed) {

package/src/react/runtimes/assistant-transport/transport-parentid.test.tsx

@@ -55,6 +55,14 @@ const setupRuntime = () => {
   return { App, fetchMock, requestBodies, runtimeRef };
 };
 
+const toolResult = (toolCallId: string): AssistantTransportCommand => ({
+  type: "add-tool-result",
+  toolCallId,
+  toolName: "tool",
+  result: {},
+  isError: false,
+});
+
 describe("assistant transport parentId lifetime", () => {
   afterEach(() => {
     vi.unstubAllGlobals();
@@ -94,4 +102,53 @@ describe("assistant transport parentId lifetime", () => {
     await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
     expect(Object.hasOwn(requestBodies[1]!, "parentId")).toBe(false);
   });
+
+  it.each(["cancel", "error"] as const)(
+    "drops the parentId of a message discarded by a run's %s",
+    async (outcome) => {
+      const { App, fetchMock, requestBodies, runtimeRef } = setupRuntime();
+      let settleFirstRun!: () => void;
+      fetchMock.mockImplementationOnce(
+        (_url: RequestInfo | URL, init?: RequestInit) => {
+          requestBodies.push(JSON.parse(init!.body as string));
+          return new Promise<Response>((_resolve, reject) => {
+            settleFirstRun = () => reject(new Error("network down"));
+            init!.signal!.addEventListener(
+              "abort",
+              () => reject(init!.signal!.reason),
+              { once: true },
+            );
+          });
+        },
+      );
+
+      await act(async () => {
+        render(<App />);
+      });
+      await waitFor(() => expect(runtimeRef.current).not.toBeNull());
+      const extras = () =>
+        runtimeRef.current!.thread.getState().extras as {
+          sendCommand: (command: AssistantTransportCommand) => void;
+        };
+
+      act(() => extras().sendCommand(toolResult("t0")));
+      await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+      await act(async () => {
+        void runtimeRef.current!.thread.append("m2");
+      });
+      await act(async () => {
+        if (outcome === "cancel") runtimeRef.current!.thread.cancelRun();
+        else settleFirstRun();
+      });
+      await waitFor(() =>
+        expect(runtimeRef.current!.thread.getState().isRunning).toBe(false),
+      );
+
+      act(() => extras().sendCommand(toolResult("t1")));
+
+      await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
+      expect(requestBodies[1]!["commands"]).toEqual([toolResult("t1")]);
+      expect(Object.hasOwn(requestBodies[1]!, "parentId")).toBe(false);
+    },
+  );
 });