@assistant-ui/react 0.15.4 → 0.15.6

package/src/legacy-runtime/AssistantRuntimeProvider.tsx

@@ -1,7 +1,11 @@
 "use client";
 
 import { type FC, memo, type PropsWithChildren, useEffect } from "react";
-import { type AssistantClient, useAui } from "@assistant-ui/store";
+import {
+  type AssistantClient,
+  type AuiConfig,
+  useAui,
+} from "@assistant-ui/store";
 import type { AssistantRuntime } from "./runtime/AssistantRuntime";
 import { AssistantProviderBase } from "@assistant-ui/core/react";
 import { ThreadPrimitiveViewportProvider } from "../context/providers/ThreadViewportProvider";
@@ -24,6 +28,12 @@ export namespace AssistantRuntimeProvider {
      * @defaultValue undefined
      */
     aui?: AssistantClient;
+
+    /**
+     * Optional extra scopes provided alongside the runtime's `threads`
+     * scope; build with `AuiConfig`.
+     */
+    config?: AuiConfig;
   }>;
 }
 
@@ -39,9 +49,9 @@ const DevToolsRegistration: FC = () => {
 
 export const AssistantRuntimeProviderImpl: FC<
   AssistantRuntimeProvider.Props
-> = ({ children, aui, runtime }) => {
+> = ({ children, aui, config, runtime }) => {
   return (
-    <AssistantProviderBase runtime={runtime} aui={aui ?? null}>
+    <AssistantProviderBase runtime={runtime} aui={aui ?? null} config={config}>
       <DevToolsRegistration />
       {/* TODO temporarily allow accessing viewport state from outside the viewport */}
       {/* TODO figure out if this behavior should be deprecated, since it is quite hacky */}

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

@@ -9,6 +9,13 @@ export type RunManager = Readonly<{
 
 const disposeReason = Symbol("assistant-transport-dispose");
 
+const reportFinishError = (error: unknown) => {
+  console.error(
+    "[assistant-ui] Assistant transport onFinish callback threw an error",
+    error,
+  );
+};
+
 export function useRunManager(config: {
   onRun: (signal: AbortSignal) => Promise<void>;
   onFinish?: (() => void) | undefined;
@@ -52,7 +59,13 @@ export function useRunManager(config: {
         }
       } finally {
         if (!disposeAborted() && !stateRef.current.disposed) {
-          onFinishRef.current?.();
+          try {
+            void Promise.resolve(onFinishRef.current?.()).catch(
+              reportFinishError,
+            );
+          } catch (error) {
+            reportFinishError(error);
+          }
         }
         if (!stateRef.current.disposed && stateRef.current.pending) {
           startRun();

package/src/legacy-runtime/runtime-cores/assistant-transport/transport-scheduling.test.ts

@@ -146,6 +146,44 @@ describe("assistant transport scheduling contracts", () => {
     expect(seen[0]).toEqual([createMessageCommand("m1")]);
   });
 
+  it("settles the run when onFinish throws", async () => {
+    const error = new Error("telemetry failed");
+    const onFinish = vi.fn<() => void>().mockImplementationOnce(() => {
+      throw error;
+    });
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    try {
+      const { result } = renderHook(() =>
+        useTransportSchedulingHarness({ onFinish }),
+      );
+
+      act(() => {
+        result.current.commandQueue.enqueue(createMessageCommand("m1"));
+      });
+
+      await waitFor(() => {
+        expect(result.current.runManager.isRunning).toBe(false);
+      });
+
+      act(() => {
+        result.current.commandQueue.enqueue(createMessageCommand("m2"));
+      });
+
+      await waitFor(() => {
+        expect(result.current.runBatchesRef.current).toHaveLength(2);
+      });
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui] Assistant transport onFinish callback threw an error",
+        error,
+      );
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
   it("cancel returns combined in-flight and queued commands", async () => {
     const onCancel = vi.fn();
     const { result } = renderHook(() =>

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

@@ -92,7 +92,9 @@ export type AssistantTransportProtocol = "data-stream" | "assistant-transport";
 
 export type SendCommandsRequestBody = {
   commands: QueuedCommand[];
-  state: unknown;
+  /** Absent on a resume with `resumeStateApi`; the server replays from its retained snapshot. */
+  state?: unknown;
+  runId?: string;
   system: string | undefined;
   tools: Record<string, unknown> | undefined;
   callSettings: LanguageModelV1CallSettings | undefined;
@@ -109,6 +111,8 @@ export type AssistantTransportOptions<T> = {
   initialState: T;
   api: string;
   resumeApi?: string;
+  /** Endpoint that returns the retained initial state and run ID for a resume stream. A 204 response means no run is active and the resume is skipped. */
+  resumeStateApi?: string;
   protocol?: AssistantTransportProtocol;
   /**
    * When `false`, stream decoding and state reconciliation tolerate malformed

package/src/legacy-runtime/runtime-cores/assistant-transport/useAssistantTransport.spec.md

@@ -16,6 +16,13 @@ Command Scheduling
   - If no run is in progress: start a run immediately and flush commands to the server.
 - A follow-up run that finds an empty queue is a no-op: no request is sent and no error is surfaced.
 - A resume run sends no commands; commands enqueued while it is pending or active are flushed in a follow-up run after it settles.
+
+Resume State
+
+- With `resumeStateApi` configured, a resume first posts `{ threadId }` there; the endpoint returns `{ runId, state }`, the state that started the active run.
+- The resume request carries `runId` and no `state`: the server replays from the snapshot it retained for that run, so neither `body` overrides nor a `prepareSendCommandsRequest` rebuild can substitute a different base. `runId` takes precedence over `body` fields and is re-attached after `prepareSendCommandsRequest` so a rebuilt body cannot drop it. The local base is replaced by the snapshot only after the resume stream responds OK, so a rejected resume keeps the local state.
+- A 204 response means no active run: the resume is skipped without error, and queued commands are flushed in a follow-up run.
+- A malformed snapshot response fails the resume before the replay request is sent.
 - Runs execute on a `queueMicrotask`, so multiple synchronous enqueues coalesce into a single request: the first run's flush takes all of them, and the coalesced follow-up run no-ops.
 
 Command Queue

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>
   );
 });