@assistant-ui/react 0.15.4 → 0.15.5

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime.test.tsx

@@ -137,6 +137,7 @@ describe("useAssistantTransportRuntime", () => {
     expect(
       fetchMock.requests[0]!.body["commands"].map((c: any) => c.type),
     ).toEqual(["add-message", "add-message"]);
+    expect(fetchMock.requests[0]!.body["state"]).toEqual({});
 
     act(() => fetchMock.servers[0]!.close());
 
@@ -201,6 +202,7 @@ describe("useAssistantTransportRuntime", () => {
     await waitFor(() => expect(fetchMock.requests).toHaveLength(2));
     expect(fetchMock.requests[1]!.url).toBe("https://example.com/resume");
     expect(fetchMock.requests[1]!.body["commands"]).toEqual([]);
+    expect(fetchMock.requests[1]!.body).toHaveProperty("state");
 
     // "b" coalesced into the resume run and must not starve in the queue.
     act(() => fetchMock.servers[1]!.close());
@@ -213,4 +215,292 @@ describe("useAssistantTransportRuntime", () => {
     act(() => fetchMock.servers[2]!.close());
     await waitFor(() => expect(aui().thread.getState().isRunning).toBe(false));
   });
+
+  it("applies resumed operations to the retained initial state", async () => {
+    const requests: RecordedRequest[] = [];
+    vi.stubGlobal(
+      "fetch",
+      async (url: RequestInfo | URL, init: RequestInit = {}) => {
+        requests.push({
+          url: String(url),
+          init,
+          body: JSON.parse(init.body as string),
+        });
+
+        if (String(url) === "https://example.com/resume-state") {
+          return Response.json({
+            runId: "run-1",
+            state: { message: "Hello" },
+          });
+        }
+
+        return new Response(
+          'aui-state:[{"type":"append-text","path":["message"],"value":" world"}]\n',
+          { status: 200 },
+        );
+      },
+    );
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    act(() => {
+      aui().thread.importExternalState({ message: "Wrong" });
+    });
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { state: unknown }).state,
+      ).toEqual({ message: "Hello world" }),
+    );
+    expect(requests.map((request) => request.url)).toEqual([
+      "https://example.com/resume-state",
+      "https://example.com/resume",
+    ]);
+    expect(requests[1]!.body).toMatchObject({ runId: "run-1" });
+    expect(requests[1]!.body).not.toHaveProperty("state");
+  });
+
+  it("rejects malformed resume state responses before replay", async () => {
+    const fetchMock = vi.fn(async () => Response.json({ state: {} }));
+    vi.stubGlobal("fetch", fetchMock);
+    const onError = vi.fn();
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+      onError,
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    await waitFor(() =>
+      expect(onError).toHaveBeenCalledWith(
+        expect.objectContaining({
+          message: "Resume state response must contain state and runId",
+        }),
+        expect.anything(),
+      ),
+    );
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("commits a retained null state locally and omits state from the resume request", async () => {
+    const requests: RecordedRequest[] = [];
+    vi.stubGlobal(
+      "fetch",
+      async (url: RequestInfo | URL, init: RequestInit = {}) => {
+        requests.push({
+          url: String(url),
+          init,
+          body: JSON.parse(init.body as string),
+        });
+
+        if (String(url) === "https://example.com/resume-state") {
+          return Response.json({ runId: "run-1", state: null });
+        }
+
+        return new Response("", { status: 200 });
+      },
+    );
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    act(() => {
+      aui().thread.importExternalState({ message: "Wrong" });
+    });
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    expect(requests[1]!.body).toMatchObject({ runId: "run-1" });
+    expect(requests[1]!.body).not.toHaveProperty("state");
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { state: unknown }).state,
+      ).toBeNull(),
+    );
+  });
+
+  it("skips the resume without error when the state endpoint reports no active run", async () => {
+    const fetchMock = vi.fn(async () => new Response(null, { status: 204 }));
+    vi.stubGlobal("fetch", fetchMock);
+    const onError = vi.fn();
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+      onError,
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    act(() => {
+      aui().thread.importExternalState({ message: "Kept" });
+    });
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    await waitFor(() => expect(aui().thread.getState().isRunning).toBe(false));
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(onError).not.toHaveBeenCalled();
+    expect(
+      (aui().thread.getState().extras as { state: unknown }).state,
+    ).toEqual({ message: "Kept" });
+  });
+
+  it("keeps the retained runId over body overrides in the resume request", async () => {
+    const requests: RecordedRequest[] = [];
+    vi.stubGlobal(
+      "fetch",
+      async (url: RequestInfo | URL, init: RequestInit = {}) => {
+        requests.push({
+          url: String(url),
+          init,
+          body: JSON.parse(init.body as string),
+        });
+
+        if (String(url) === "https://example.com/resume-state") {
+          return Response.json({
+            runId: "run-1",
+            state: { message: "Hello" },
+          });
+        }
+
+        return new Response("", { status: 200 });
+      },
+    );
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+      body: { state: { message: "Injected" }, runId: "bogus" },
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    expect(requests[1]!.body["runId"]).toBe("run-1");
+    expect(requests[1]!.body).not.toHaveProperty("state");
+  });
+
+  it("re-attaches runId and strips substituted state when prepareSendCommandsRequest rebuilds the body", async () => {
+    const requests: RecordedRequest[] = [];
+    vi.stubGlobal(
+      "fetch",
+      async (url: RequestInfo | URL, init: RequestInit = {}) => {
+        requests.push({
+          url: String(url),
+          init,
+          body: JSON.parse(init.body as string),
+        });
+
+        if (String(url) === "https://example.com/resume-state") {
+          return Response.json({
+            runId: "run-1",
+            state: { message: "Hello" },
+          });
+        }
+
+        return new Response("", { status: 200 });
+      },
+    );
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+      prepareSendCommandsRequest: (body) => ({
+        commands: body.commands,
+        state: { message: "Substituted" },
+        rebuilt: true,
+      }),
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    expect(requests[1]!.body).toMatchObject({ runId: "run-1", rebuilt: true });
+    expect(requests[1]!.body).not.toHaveProperty("state");
+  });
+
+  it("keeps local state when the matching resume stream is rejected", async () => {
+    const fetchMock = vi
+      .fn()
+      .mockResolvedValueOnce(
+        Response.json({ runId: "run-1", state: { message: "Hello" } }),
+      )
+      .mockResolvedValueOnce(new Response("run mismatch", { status: 409 }));
+    vi.stubGlobal("fetch", fetchMock);
+    const onError = vi.fn();
+    const { aui } = mountRuntime({
+      resumeApi: "https://example.com/resume",
+      resumeStateApi: "https://example.com/resume-state",
+      onError,
+    });
+    await waitFor(() =>
+      expect(
+        (aui().thread.getState().extras as { sendCommand?: unknown })
+          ?.sendCommand,
+      ).toBeTypeOf("function"),
+    );
+
+    act(() => {
+      aui().thread.importExternalState({ message: "Wrong" });
+    });
+    await act(async () => {
+      await aui().thread.resumeRun({ parentId: null });
+    });
+
+    await waitFor(() =>
+      expect(onError).toHaveBeenCalledWith(
+        expect.objectContaining({ message: "Status 409: run mismatch" }),
+        expect.anything(),
+      ),
+    );
+    expect(
+      (aui().thread.getState().extras as { state: unknown }).state,
+    ).toEqual({ message: "Wrong" });
+    expect(fetchMock).toHaveBeenCalledTimes(2);
+  });
 });

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransportRuntime.ts

@@ -72,6 +72,36 @@ const convertAppendMessageToCommand = (
   };
 };
 
+const readResumeState = async <T>(
+  response: Response,
+): Promise<{ runId: string; state: T } | null> => {
+  if (response.status === 204) return null;
+  if (!response.ok) {
+    throw new Error(
+      `Resume state request failed with status ${response.status}: ${await response.text()}`,
+    );
+  }
+
+  let value: unknown;
+  try {
+    value = await response.json();
+  } catch {
+    throw new Error("Resume state response was not valid JSON");
+  }
+
+  if (
+    typeof value !== "object" ||
+    value === null ||
+    !("state" in value) ||
+    !("runId" in value) ||
+    typeof value.runId !== "string"
+  ) {
+    throw new Error("Resume state response must contain state and runId");
+  }
+
+  return { runId: value.runId, state: value.state as T };
+};
+
 const symbolAssistantTransportExtras = Symbol("assistant-transport-extras");
 type AssistantTransportExtras = {
   [symbolAssistantTransportExtras]: true;
@@ -160,6 +190,24 @@ const useAssistantTransportThreadRuntime = <T>(
       if (!isResume) parentIdRef.current = undefined;
 
       const headers = await createRequestHeaders(options.headers);
+      let resumeState: { runId: string; state: T } | undefined;
+      if (isResume && options.resumeStateApi) {
+        const resumeStateResponse = await fetch(options.resumeStateApi, {
+          method: "POST",
+          headers,
+          body: JSON.stringify({ threadId }),
+          signal,
+        });
+        const retained = await readResumeState<T>(resumeStateResponse);
+        if (retained === null) {
+          if (commandQueue.state.queued.length > 0) {
+            runManager.schedule();
+          }
+          return;
+        }
+        resumeState = retained;
+      }
+
       const bodyValue =
         typeof options.body === "function"
           ? await options.body()
@@ -168,7 +216,7 @@ const useAssistantTransportThreadRuntime = <T>(
 
       let requestBody: Record<string, unknown> = {
         commands,
-        state: agentStateRef.current,
+        ...(resumeState === undefined && { state: agentStateRef.current }),
         system: context.system,
         tools: context.tools ? toToolsJSONSchema(context.tools) : undefined,
         threadId,
@@ -190,6 +238,14 @@ const useAssistantTransportThreadRuntime = <T>(
         );
       }
 
+      if (resumeState !== undefined) {
+        // The server replays a resume from the snapshot it retained for this
+        // runId. Body overrides and prepare hooks can neither substitute a
+        // state nor drop the ID the server validates against.
+        requestBody = { ...requestBody, runId: resumeState.runId };
+        delete requestBody["state"];
+      }
+
       const response = await fetch(
         isResume ? options.resumeApi! : options.api,
         {
@@ -210,6 +266,11 @@ const useAssistantTransportThreadRuntime = <T>(
         throw new Error("Response body is null");
       }
 
+      if (resumeState !== undefined) {
+        agentStateRef.current = resumeState.state;
+        rerender((prev) => prev + 1);
+      }
+
       const body = await createReplayBoundaryStream(response, {
         setReplaying: setIsReplaying,
         waitForRender: waitForReplayRender,

package/src/primitives/actionBar/ActionBarExportMarkdown.tsx

@@ -37,7 +37,7 @@ const useActionBarExportMarkdown = ({
     a.href = url;
     a.download = filename ?? `message-${Date.now()}.md`;
     a.click();
-    URL.revokeObjectURL(url);
+    setTimeout(() => URL.revokeObjectURL(url), 40_000);
   }, [aui, filename, onExport]);
 
   if (!hasExportableContent) return null;

package/src/primitives/attachment/AttachmentThumb.test.tsx

@@ -0,0 +1,70 @@
+import { describe, expect, it, vi } from "vitest";
+import { renderToStaticMarkup } from "react-dom/server";
+import type * as AssistantStore from "@assistant-ui/store";
+import { AttachmentPrimitiveThumb } from "./AttachmentThumb";
+
+const mockUseAuiState = vi.fn();
+type UseAuiStateSelector = Parameters<
+  (typeof AssistantStore)["useAuiState"]
+>[0];
+
+vi.mock("@assistant-ui/store", async (importOriginal) => {
+  const actual = await importOriginal<typeof AssistantStore>();
+  return {
+    ...actual,
+    useAuiState: (selector: UseAuiStateSelector) => mockUseAuiState(selector),
+  };
+});
+
+const renderThumb = (
+  name: string,
+  type = "file",
+  props?: AttachmentPrimitiveThumb.Props,
+) => {
+  mockUseAuiState.mockImplementation((selector: UseAuiStateSelector) =>
+    selector({ attachment: { name, type } } as never),
+  );
+
+  return renderToStaticMarkup(<AttachmentPrimitiveThumb {...props} />);
+};
+
+describe("AttachmentPrimitiveThumb", () => {
+  it("renders the dotted extension for a single-extension name", () => {
+    expect(renderThumb("photo.png")).toBe("<div>.png</div>");
+  });
+
+  it("takes only the last segment for a multi-extension name", () => {
+    expect(renderThumb("archive.tar.gz")).toBe("<div>.gz</div>");
+  });
+
+  it("falls back to the attachment type when the name has no extension", () => {
+    expect(renderThumb("noext", "document")).toBe("<div>document</div>");
+  });
+
+  it("falls back to the attachment type for an empty name", () => {
+    expect(renderThumb("", "image")).toBe("<div>image</div>");
+  });
+
+  it("treats a leading-dot name as extensionless", () => {
+    expect(renderThumb(".gitignore")).toBe("<div>file</div>");
+  });
+
+  it("falls back to the attachment type for a trailing-dot name", () => {
+    expect(renderThumb("report.", "image")).toBe("<div>image</div>");
+  });
+
+  it("renders custom children instead of the label", () => {
+    expect(renderThumb("report.pdf", "file", { children: <em>PDF</em> })).toBe(
+      "<div><em>PDF</em></div>",
+    );
+  });
+
+  it("renders the child element when asChild is set", () => {
+    expect(
+      renderThumb("report.pdf", "file", {
+        asChild: true,
+        children: <span>custom</span>,
+      }),
+    ).toBe("<span>custom</span>");
+  });
+});

package/src/primitives/attachment/AttachmentThumb.tsx

@@ -19,13 +19,17 @@ export const AttachmentPrimitiveThumb = forwardRef<
   AttachmentPrimitiveThumb.Element,
   AttachmentPrimitiveThumb.Props
 >((props, ref) => {
-  const ext = useAuiState((s) => {
-    const parts = s.attachment.name.split(".");
-    return parts.length > 1 ? parts.pop()! : "";
+  const label = useAuiState((s) => {
+    const name = s.attachment.name;
+    const dot = name.lastIndexOf(".");
+    if (dot > 0 && dot < name.length - 1) {
+      return `.${name.slice(dot + 1)}`;
+    }
+    return s.attachment.type;
   });
   return (
     <Primitive.div {...props} ref={ref}>
-      .{ext}
+      {props.children ?? label}
     </Primitive.div>
   );
 });

package/src/primitives/composer/ComposerInput.test.tsx

@@ -10,6 +10,7 @@ const setText = vi.fn<(text: string) => void>();
 const setCursorPosition = vi.fn<(pos: number) => void>();
 const sendSpy = vi.fn<(options?: { steer?: boolean }) => void>();
 const addAttachment = vi.fn<(file: File) => Promise<void>>();
+const cancelSpy = vi.fn<() => void>();
 
 const composerState = {
   isEditing: true,
@@ -41,7 +42,7 @@ vi.mock("@assistant-ui/store", () => {
     composer: {
       setText: (text: string) => setText(text),
       getState: () => composerState,
-      cancel: () => {},
+      cancel: () => cancelSpy(),
       send: (options?: { steer?: boolean }) => sendSpy(options),
       addAttachment: (file: File) => addAttachment(file),
     },
@@ -78,8 +79,12 @@ vi.mock("./trigger/TriggerPopoverRootContext", () => ({
   useTriggerPopoverActiveAriaOptional: () => activeAria,
 }));
 
+let escapeKeydownHandler: ((event: KeyboardEvent) => void) | null = null;
+
 vi.mock("@radix-ui/react-use-escape-keydown", () => ({
-  useEscapeKeydown: () => {},
+  useEscapeKeydown: (handler: (event: KeyboardEvent) => void) => {
+    escapeKeydownHandler = handler;
+  },
 }));
 
 vi.mock("../../utils/hooks/useOnScrollToBottom", () => ({
@@ -181,16 +186,19 @@ describe("ComposerPrimitiveInput", () => {
     setCursorPosition.mockReset();
     sendSpy.mockReset();
     addAttachment.mockReset();
+    cancelSpy.mockReset();
     composerState.isEditing = true;
     composerState.text = "";
     composerState.isEmpty = true;
     composerState.canSend = true;
     composerState.dictation = undefined;
+    composerState.canCancel = false;
     threadState.isDisabled = false;
     threadState.isRunning = false;
     threadState.capabilities = { queue: false, attachments: false };
     pluginRegistry = null;
     activeAria = null;
+    escapeKeydownHandler = null;
     setMatchMedia(false);
 
     requestSubmitSpy = vi.fn<(submitter?: HTMLElement | null) => void>();
@@ -514,6 +522,49 @@ describe("ComposerPrimitiveInput", () => {
     });
   });
 
+  describe("escape behavior", () => {
+    const fireEscape = (
+      textarea: HTMLTextAreaElement,
+      opts: { isComposing?: boolean } = {},
+    ): KeyboardEvent => {
+      const event = new KeyboardEvent("keydown", {
+        bubbles: true,
+        cancelable: true,
+        key: "Escape",
+        isComposing: opts.isComposing ?? false,
+      });
+      textarea.dispatchEvent(event);
+      escapeKeydownHandler?.(event);
+      return event;
+    };
+
+    it("cancels the composer on Escape when cancellable", async () => {
+      composerState.canCancel = true;
+      const textarea = await mount();
+
+      let event!: KeyboardEvent;
+      await act(async () => {
+        event = fireEscape(textarea);
+      });
+
+      expect(cancelSpy).toHaveBeenCalledTimes(1);
+      expect(event.defaultPrevented).toBe(true);
+    });
+
+    it("ignores Escape while an IME composition is active", async () => {
+      composerState.canCancel = true;
+      const textarea = await mount();
+
+      let event!: KeyboardEvent;
+      await act(async () => {
+        event = fireEscape(textarea, { isComposing: true });
+      });
+
+      expect(cancelSpy).not.toHaveBeenCalled();
+      expect(event.defaultPrevented).toBe(false);
+    });
+  });
+
   describe("paste attachments", () => {
     it("attempts every pasted file and prevents default when attachments are supported", async () => {
       threadState.capabilities = { queue: false, attachments: true };

package/src/primitives/composer/ComposerInput.tsx

@@ -199,6 +199,9 @@ export const ComposerPrimitiveInput = forwardRef<
       // Only handle ESC if it originated from within this input
       if (!textareaRef.current?.contains(e.target as Node)) return;
 
+      // ignore IME composition events
+      if (e.isComposing) return;
+
       // Let registered plugins (mention, slash command, etc.) handle Escape first
       if (pluginRegistry) {
         for (const plugin of pluginRegistry.getPlugins()) {

package/src/primitives/queueItem/QueueItemText.tsx

@@ -25,11 +25,17 @@ export const QueueItemPrimitiveText = forwardRef<
   QueueItemPrimitiveText.Element,
   QueueItemPrimitiveText.Props
 >((props, ref) => {
-  const prompt = useAuiState((s) => s.queueItem.prompt);
+  const text = useAuiState((s) =>
+    // hosts on the pre-parts adapter shape may omit the field at runtime
+    (s.queueItem.parts ?? [])
+      .filter((part) => part.type === "text")
+      .map((part) => part.text)
+      .join("\n\n"),
+  );
 
   return (
     <Primitive.span {...props} ref={ref}>
-      {props.children ?? prompt}
+      {props.children ?? text}
     </Primitive.span>
   );
 });