@assistant-ui/react-ai-sdk 1.3.23 → 1.3.26

package/src/ui/use-chat/AssistantChatTransport.ts

@@ -8,9 +8,22 @@ import {
   type UIMessage,
 } from "ai";
 import { toToolsJSONSchema } from "assistant-stream";
+import {
+  RESUMABLE_STREAM_ID_HEADER,
+  type AssistantChatResumableOptions,
+} from "../resumable";
 
 type InitializableThreadListItem = Pick<ThreadListItemRuntime, "initialize">;
 
+const FINISH_MARKER = '"type":"finish"';
+const FINISH_BUFFER_LIMIT = 4096;
+const FINISH_BUFFER_TAIL = 1024;
+
+export type AssistantChatTransportInitOptions<UI_MESSAGE extends UIMessage> =
+  HttpChatTransportInitOptions<UI_MESSAGE> & {
+    resumable?: AssistantChatResumableOptions;
+  };
+
 export class AssistantChatTransport<
   UI_MESSAGE extends UIMessage,
 > extends DefaultChatTransport<UI_MESSAGE> {
@@ -18,9 +31,22 @@ export class AssistantChatTransport<
   private getThreadListItem:
     | (() => InitializableThreadListItem | undefined)
     | undefined;
-  constructor(initOptions?: HttpChatTransportInitOptions<UI_MESSAGE>) {
+  private readonly resumable: AssistantChatResumableOptions | undefined;
+
+  constructor(initOptions?: AssistantChatTransportInitOptions<UI_MESSAGE>) {
+    const { resumable, ...rest } = initOptions ?? {};
+    const userFetch = rest.fetch;
+    const userPrepareReconnect = rest.prepareReconnectToStreamRequest;
+
     super({
-      ...initOptions,
+      ...rest,
+      ...(resumable && {
+        fetch: wrapFetchWithResumable(resumable, userFetch),
+        prepareReconnectToStreamRequest: wrapPrepareReconnect(
+          resumable,
+          userPrepareReconnect,
+        ),
+      }),
       prepareSendMessagesRequest: async (options) => {
         const context = this.runtime?.thread.getModelContext();
         const threadListItem =
@@ -38,7 +64,7 @@ export class AssistantChatTransport<
           },
         };
         const preparedRequest =
-          await initOptions?.prepareSendMessagesRequest?.(optionsEx);
+          await rest.prepareSendMessagesRequest?.(optionsEx);
 
         return {
           ...preparedRequest,
@@ -53,15 +79,90 @@ export class AssistantChatTransport<
         };
       },
     });
+
+    this.resumable = resumable;
   }
 
   setRuntime(runtime: AssistantRuntime) {
     this.runtime = runtime;
   }
 
+  getResumableAdapter(): AssistantChatResumableOptions | undefined {
+    return this.resumable;
+  }
+
   __internal_setGetThreadListItem(
     getter: () => InitializableThreadListItem | undefined,
   ) {
     this.getThreadListItem = getter;
   }
 }
+
+function wrapFetchWithResumable(
+  resumable: AssistantChatResumableOptions,
+  userFetch: HttpChatTransportInitOptions<UIMessage>["fetch"],
+): NonNullable<HttpChatTransportInitOptions<UIMessage>["fetch"]> {
+  const baseFetch: typeof globalThis.fetch = userFetch
+    ? (input, init) => userFetch(input as RequestInfo | URL, init)
+    : globalThis.fetch.bind(globalThis);
+
+  return async (input, init) => {
+    const res = await baseFetch(input, init);
+    const id = res.headers.get(RESUMABLE_STREAM_ID_HEADER);
+    if (id) resumable.storage.setStreamId(id);
+    if (!res.body) return res;
+
+    const detectFinish = resumable.isFinishEvent ?? defaultIsFinishEvent;
+    // a single decoder is required so multi-byte sequences split across
+    // chunks buffer via stream: true rather than getting dropped.
+    const decoder = new TextDecoder();
+    let accumulator = "";
+    const tap = new TransformStream<Uint8Array, Uint8Array>({
+      transform(chunk, controller) {
+        controller.enqueue(chunk);
+        accumulator += decoder.decode(chunk, { stream: true });
+        if (detectFinish(chunk, accumulator)) {
+          resumable.storage.clear();
+          accumulator = "";
+        } else if (accumulator.length > FINISH_BUFFER_LIMIT) {
+          accumulator = accumulator.slice(-FINISH_BUFFER_TAIL);
+        }
+      },
+    });
+
+    return new Response(res.body.pipeThrough(tap), {
+      status: res.status,
+      statusText: res.statusText,
+      headers: res.headers,
+    });
+  };
+}
+
+function defaultIsFinishEvent(_chunk: Uint8Array, accumulator: string) {
+  return accumulator.includes(FINISH_MARKER);
+}
+
+function wrapPrepareReconnect(
+  resumable: AssistantChatResumableOptions,
+  userPrepareReconnect: HttpChatTransportInitOptions<UIMessage>["prepareReconnectToStreamRequest"],
+): NonNullable<
+  HttpChatTransportInitOptions<UIMessage>["prepareReconnectToStreamRequest"]
+> {
+  return async (options) => {
+    const streamId = resumable.storage.getStreamId();
+    if (!streamId) {
+      throw new Error(
+        "AssistantChatTransport: no resumable stream id available; nothing to resume",
+      );
+    }
+    const api =
+      typeof resumable.resumeApi === "function"
+        ? resumable.resumeApi(streamId)
+        : resumable.resumeApi;
+    const userPrepared = await userPrepareReconnect?.({ ...options, api });
+    return {
+      ...userPrepared,
+      api: userPrepared?.api ?? api,
+    };
+  };
+}

package/src/ui/use-chat/useAISDKRuntime.test.ts

@@ -266,6 +266,56 @@ describe("useAISDKRuntime", () => {
     );
   });
 
+  it("forwards onResume so runtime.thread.resumeRun is delivered to the adapter", async () => {
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "first" }] },
+    ]);
+    const onResume = vi.fn().mockResolvedValue(undefined);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat, { onResume }));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages.length).toBe(1);
+    });
+
+    act(() => {
+      result.current.thread.resumeRun({
+        parentId: "u1",
+        runConfig: { custom: { turnId: "t-42" } },
+      });
+    });
+
+    await waitFor(() => {
+      expect(onResume).toHaveBeenCalledTimes(1);
+    });
+
+    expect(onResume).toHaveBeenCalledWith(
+      expect.objectContaining({
+        parentId: "u1",
+        sourceId: null,
+        runConfig: { custom: { turnId: "t-42" } },
+      }),
+    );
+  });
+
+  it("rejects when resumeRun is called without an onResume adapter", async () => {
+    const chat = createChatHelpers([
+      { id: "u1", role: "user", parts: [{ type: "text", text: "first" }] },
+    ]);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages.length).toBe(1);
+    });
+
+    await expect(
+      result.current.thread.resumeRun({
+        parentId: "u1",
+      }) as unknown as Promise<void>,
+    ).rejects.toThrow("Runtime does not support resuming runs.");
+  });
+
   it("reload slices history and regenerates with metadata", async () => {
     const chat = createChatHelpers([
       { id: "u1", role: "user", parts: [{ type: "text", text: "first" }] },

package/src/ui/use-chat/useAISDKRuntime.ts

@@ -14,18 +14,22 @@ import type {
   ThreadHistoryAdapter,
   AssistantRuntime,
   ThreadMessage,
+  ThreadSuggestion,
   MessageFormatAdapter,
   MessageFormatItem,
   MessageFormatRepository,
   AppendMessage,
   RunConfig,
+  McpAppMetadata,
 } from "@assistant-ui/core";
 import { getExternalStoreMessages } from "@assistant-ui/core";
+import type { ReadonlyJSONObject } from "assistant-stream/utils";
 import { sliceMessagesUntil } from "../utils/sliceMessagesUntil";
 import { toCreateMessage } from "../utils/toCreateMessage";
 import { vercelAttachmentAdapter } from "../utils/vercelAttachmentAdapter";
 import { getVercelAIMessages } from "../getVercelAIMessages";
 import { AISDKMessageConverter } from "../utils/convertMessage";
+import { wrapModelContentEnvelope } from "../../modelContentEnvelope";
 import {
   type AISDKStorageFormat,
   aiSDKV6FormatAdapter,
@@ -68,6 +72,21 @@ export type AISDKRuntimeAdapter = {
    * @default true
    */
   cancelPendingToolCallsOnSend?: boolean | undefined;
+  /**
+   * Called when `runtime.thread.resumeRun(config)` is invoked.
+   *
+   * When omitted, `resumeRun` throws `"Runtime does not support resuming runs."`.
+   * Provide this to bridge resume invocations into a custom replay channel
+   * (for example, an SSE reconnect endpoint keyed by turn id).
+   */
+  onResume?: ExternalStoreAdapter["onResume"];
+  /**
+   * Follow up suggestions to surface on the thread. Use this to drive
+   * dynamic suggestions from application state, tool results, or backend
+   * responses; flows into `thread.suggestions` and is rendered by
+   * components that read it (such as the shadcn `ThreadFollowupSuggestions`).
+   */
+  suggestions?: readonly ThreadSuggestion[] | undefined;
 };
 
 export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
@@ -76,6 +95,8 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     adapters,
     toCreateMessage: customToCreateMessage,
     cancelPendingToolCallsOnSend = true,
+    onResume,
+    suggestions,
   }: AISDKRuntimeAdapter = {},
 ) => {
   const contextAdapters = useRuntimeAdapters();
@@ -85,6 +106,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
     new Map(),
   );
+  const toolLastInputCacheRef = useRef<Map<string, ReadonlyJSONObject>>(
+    new Map(),
+  );
+  const mcpAppMetadataCacheRef = useRef<Map<string, McpAppMetadata>>(new Map());
   const lastRunConfigRef = useRef<RunConfig | undefined>(undefined);
 
   const hasExecutingTools = Object.values(toolStatuses).some(
@@ -105,6 +130,8 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         toolStatuses,
         messageTiming,
         toolArgsKeyOrderCache: toolArgsKeyOrderCacheRef.current,
+        toolLastInputCache: toolLastInputCacheRef.current,
+        mcpAppMetadataCache: mcpAppMetadataCacheRef.current,
         ...(chatHelpers.error && { error: chatHelpers.error.message }),
       }),
       [toolStatuses, messageTiming, chatHelpers.error],
@@ -125,10 +152,14 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     getTools: () => runtimeRef.current.thread.getModelContext().tools,
     onResult: (command) => {
       if (command.type === "add-tool-result") {
+        const output =
+          command.modelContent !== undefined
+            ? wrapModelContentEnvelope(command.result, command.modelContent)
+            : command.result;
         chatHelpers.addToolResult({
           tool: command.toolName,
           toolCallId: command.toolCallId,
-          output: command.result,
+          output,
           options: { metadata: lastRunConfigRef.current },
         });
       }
@@ -319,6 +350,8 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     },
     onResumeToolCall: (options) =>
       toolInvocations.resume(options.toolCallId, options.payload),
+    ...(onResume && { onResume }),
+    ...(suggestions && { suggestions }),
     adapters: {
       attachments: vercelAttachmentAdapter,
       ...contextAdapters,

package/src/ui/use-chat/useChatRuntime.ts

@@ -15,6 +15,7 @@ import {
 } from "./useAISDKRuntime";
 import type { ChatInit, ChatTransport } from "ai";
 import { AssistantChatTransport } from "./AssistantChatTransport";
+import type { AssistantChatResumableOptions } from "../resumable";
 import { useEffect, useMemo, useRef } from "react";
 
 export type UseChatRuntimeOptions<UI_MESSAGE extends UIMessage = UIMessage> =
@@ -22,6 +23,8 @@ export type UseChatRuntimeOptions<UI_MESSAGE extends UIMessage = UIMessage> =
     cloud?: AssistantCloud | undefined;
     adapters?: AISDKRuntimeAdapter["adapters"] | undefined;
     toCreateMessage?: CustomToCreateMessageFunction;
+    onResume?: AISDKRuntimeAdapter["onResume"];
+    suggestions?: AISDKRuntimeAdapter["suggestions"];
   };
 
 const useDynamicChatTransport = <UI_MESSAGE extends UIMessage = UIMessage>(
@@ -50,6 +53,18 @@ const useDynamicChatTransport = <UI_MESSAGE extends UIMessage = UIMessage>(
   return dynamicTransport;
 };
 
+const getResumableAdapter = <UI_MESSAGE extends UIMessage>(
+  transport: ChatTransport<UI_MESSAGE>,
+): AssistantChatResumableOptions | undefined => {
+  if (transport instanceof AssistantChatTransport) {
+    return transport.getResumableAdapter();
+  }
+  const candidate = (transport as { getResumableAdapter?: () => unknown })
+    .getResumableAdapter;
+  if (typeof candidate !== "function") return undefined;
+  return candidate.call(transport) as AssistantChatResumableOptions | undefined;
+};
+
 const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   options?: UseChatRuntimeOptions<UI_MESSAGE>,
 ): AssistantRuntime => {
@@ -57,6 +72,8 @@ const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     adapters,
     transport: transportOptions,
     toCreateMessage,
+    onResume,
+    suggestions,
     ...chatOptions
   } = options ?? {};
 
@@ -80,6 +97,8 @@ const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const runtime = useAISDKRuntime(chat, {
     adapters,
     ...(toCreateMessage && { toCreateMessage }),
+    ...(onResume && { onResume }),
+    ...(suggestions && { suggestions }),
   });
 
   if (transport instanceof AssistantChatTransport) {
@@ -89,6 +108,25 @@ const useChatThreadRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     );
   }
 
+  // biome-ignore lint/correctness/useHookAtTopLevel: intentional conditional/nested hook usage
+  const resumeFiredRef = useRef(false);
+  // biome-ignore lint/correctness/useHookAtTopLevel: intentional conditional/nested hook usage
+  useEffect(() => {
+    if (resumeFiredRef.current) return;
+    const adapter = getResumableAdapter(transport);
+    if (!adapter) return;
+    const pending = adapter.storage.getStreamId();
+    if (!pending) return;
+    resumeFiredRef.current = true;
+    chat.resumeStream().catch((err: unknown) => {
+      console.warn(
+        "[assistant-ui] resumable: resume failed; clearing stored stream id",
+        err,
+      );
+      adapter.storage.clear();
+    });
+  }, [transport, chat]);
+
   return runtime;
 };