@assistant-ui/ai-sdk 0.0.1

package/src/transport/AssistantChatTransport.test.ts

@@ -0,0 +1,300 @@
+import type { UIMessage } from "ai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import { RESUMABLE_STREAM_ID_HEADER } from "./resumable";
+import { AssistantChatTransport } from "./AssistantChatTransport";
+
+const emptyStreamResponse = () =>
+  new Response(
+    new ReadableStream({ start: (controller) => controller.close() }),
+    {
+      status: 200,
+      headers: { "content-type": "text/event-stream" },
+    },
+  );
+
+const createThreadListItem = (remoteId: string) => ({
+  initialize: vi.fn(async () => ({ remoteId, externalId: undefined })),
+});
+
+const sendMessagesOptions = {
+  trigger: "submit-message" as const,
+  chatId: "local-chat-id",
+  messageId: undefined,
+  messages: [{ id: "m1", role: "user", parts: [{ type: "text", text: "hi" }] }],
+  abortSignal: undefined,
+};
+
+describe("AssistantChatTransport.prepareSendMessagesRequest", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("passes the initialized remote thread id to prepareSendMessagesRequest as options.id", async () => {
+    const threadListItem = createThreadListItem("remote-thread-id");
+    const captured: { id: unknown } = { id: undefined };
+    const fetchMock = vi.fn(async () => emptyStreamResponse());
+
+    const transport = new AssistantChatTransport({
+      fetch: fetchMock as never,
+      prepareSendMessagesRequest: async (options) => {
+        captured.id = options.id;
+        return { body: { id: options.id } };
+      },
+    });
+    transport.__internal_setGetThreadListItem(() => threadListItem as never);
+
+    await transport.sendMessages(sendMessagesOptions as never);
+
+    expect(captured.id).toBe("remote-thread-id");
+    expect(threadListItem.initialize).toHaveBeenCalledTimes(1);
+  });
+
+  it("uses the initialized remote thread id in the default request body", async () => {
+    const threadListItem = createThreadListItem("remote-thread-id");
+    let capturedBody: string | undefined;
+    const fetchMock = vi.fn(
+      async (_input: unknown, init: { body?: string } | undefined) => {
+        capturedBody = init?.body;
+        return emptyStreamResponse();
+      },
+    );
+
+    const transport = new AssistantChatTransport({
+      fetch: fetchMock as never,
+    });
+    transport.__internal_setGetThreadListItem(() => threadListItem as never);
+
+    await transport.sendMessages(sendMessagesOptions as never);
+
+    expect(threadListItem.initialize).toHaveBeenCalledTimes(1);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    const body = JSON.parse(capturedBody as string);
+    expect(body.id).toBe("remote-thread-id");
+    expect(body.messages).toEqual(sendMessagesOptions.messages);
+    expect(body.trigger).toBe("submit-message");
+    expect(body.messageId).toBeUndefined();
+  });
+
+  it("falls back to the local chat id when no thread list item is available", async () => {
+    const captured: { id: unknown } = { id: undefined };
+    const fetchMock = vi.fn(async () => emptyStreamResponse());
+
+    const transport = new AssistantChatTransport({
+      fetch: fetchMock as never,
+      prepareSendMessagesRequest: async (options) => {
+        captured.id = options.id;
+        return { body: { id: options.id } };
+      },
+    });
+    transport.__internal_setGetThreadListItem(() => undefined);
+
+    await transport.sendMessages(sendMessagesOptions as never);
+
+    expect(captured.id).toBe("local-chat-id");
+  });
+
+  it("stores response stream ids under the local thread id", async () => {
+    const setStreamId = vi.fn();
+    const fetchMock = vi.fn(async (_input: string, _init?: RequestInit) => {
+      return new Response(
+        new ReadableStream({ start: (controller) => controller.close() }),
+        {
+          status: 200,
+          headers: {
+            "content-type": "text/event-stream",
+            [RESUMABLE_STREAM_ID_HEADER]: "stream-1",
+          },
+        },
+      );
+    });
+    const transport = new AssistantChatTransport({
+      fetch: fetchMock as never,
+      resumable: {
+        storage: {
+          getStreamId: vi.fn(),
+          setStreamId,
+          clear: vi.fn(),
+        },
+        resumeApi: "/api/chat/resume",
+      },
+    });
+
+    await transport.sendMessages(sendMessagesOptions as never);
+
+    expect(setStreamId).toHaveBeenCalledWith("stream-1", "local-chat-id");
+    const requestHeaders = new Headers(fetchMock.mock.calls[0]?.[1]?.headers);
+    expect(requestHeaders.has("x-assistant-ui-resumable-thread-id")).toBe(
+      false,
+    );
+  });
+
+  it("does not add the resumable thread header without resumable storage", async () => {
+    const fetchMock = vi.fn(async (_input: string, _init?: RequestInit) =>
+      emptyStreamResponse(),
+    );
+    const transport = new AssistantChatTransport({
+      fetch: fetchMock as never,
+    });
+
+    await transport.sendMessages(sendMessagesOptions as never);
+
+    const requestHeaders = new Headers(fetchMock.mock.calls[0]?.[1]?.headers);
+    expect(requestHeaders.has("x-assistant-ui-resumable-thread-id")).toBe(
+      false,
+    );
+  });
+
+  it("reads reconnect stream ids under the local thread id", async () => {
+    const getStreamId = vi.fn(() => "stream-1");
+    const fetchMock = vi.fn(async (_input: string, _init?: RequestInit) =>
+      emptyStreamResponse(),
+    );
+    const transport = new AssistantChatTransport({
+      fetch: fetchMock as never,
+      resumable: {
+        storage: {
+          getStreamId,
+          setStreamId: vi.fn(),
+          clear: vi.fn(),
+        },
+        resumeApi: (streamId) => `/api/chat/resume/${streamId}`,
+      },
+    });
+
+    await transport.reconnectToStream({ chatId: "local-chat-id" });
+
+    expect(getStreamId).toHaveBeenCalledWith("local-chat-id");
+    expect(fetchMock.mock.calls[0]?.[0]).toBe("/api/chat/resume/stream-1");
+    const requestHeaders = new Headers(fetchMock.mock.calls[0]?.[1]?.headers);
+    expect(requestHeaders.has("x-assistant-ui-resumable-thread-id")).toBe(
+      false,
+    );
+  });
+});
+
+const createMemoryStorage = (initialId: string | null = null) => {
+  let id = initialId;
+  return {
+    getStreamId: () => id,
+    setStreamId: (next: string) => {
+      id = next;
+    },
+    clear: () => {
+      id = null;
+    },
+  };
+};
+
+// A null-body status response carrying a non-null empty body, as WebKit returns
+// for a 204. The JS `Response` constructor cannot represent a null-body status
+// with a body, so this is a plain object shaped like a Response.
+const nullBodyStatusWithBody = (status: number) =>
+  ({
+    status,
+    statusText: "",
+    headers: new Headers(),
+    body: new ReadableStream({ start: (controller) => controller.close() }),
+  }) as unknown as Response;
+
+// The resumable wrapper replaces the transport's `fetch`, which is protected on
+// the upstream `HttpChatTransport`; reach it through a structural cast.
+const wrappedFetchOf = (
+  transport: AssistantChatTransport<UIMessage>,
+): ((input: string, init?: unknown) => Promise<Response>) =>
+  (
+    transport as unknown as {
+      fetch: (input: string, init?: unknown) => Promise<Response>;
+    }
+  ).fetch;
+
+describe("AssistantChatTransport resumable fetch wrapper", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("passes a 204 with a non-null empty body through untouched (WebKit)", async () => {
+    const response = nullBodyStatusWithBody(204);
+    const fetchMock = vi.fn(async () => response);
+    const transport = new AssistantChatTransport({
+      resumable: {
+        storage: createMemoryStorage(),
+        resumeApi: "/api/resume",
+      },
+      fetch: fetchMock as never,
+    });
+
+    const res = await wrappedFetchOf(transport)("https://example.com", {});
+
+    expect(res).toBe(response);
+    expect(res.status).toBe(204);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("passes a 204 with a null body through untouched", async () => {
+    const response = new Response(null, { status: 204 });
+    const fetchMock = vi.fn(async () => response);
+    const transport = new AssistantChatTransport({
+      resumable: {
+        storage: createMemoryStorage(),
+        resumeApi: "/api/resume",
+      },
+      fetch: fetchMock as never,
+    });
+
+    const res = await wrappedFetchOf(transport)("https://example.com", {});
+
+    expect(res).toBe(response);
+    expect(res.status).toBe(204);
+  });
+
+  it("passes other null-body statuses with a body through untouched", async () => {
+    const response = nullBodyStatusWithBody(304);
+    const fetchMock = vi.fn(async () => response);
+    const transport = new AssistantChatTransport({
+      resumable: {
+        storage: createMemoryStorage(),
+        resumeApi: "/api/resume",
+      },
+      fetch: fetchMock as never,
+    });
+
+    const res = await wrappedFetchOf(transport)("https://example.com", {});
+
+    expect(res).toBe(response);
+    expect(res.status).toBe(304);
+  });
+
+  it("still reconstructs and taps a normal streaming response, clearing storage on finish", async () => {
+    const storage = createMemoryStorage("stream-id-123");
+    const body = new ReadableStream({
+      start(controller) {
+        controller.enqueue(
+          new TextEncoder().encode('data: {"type":"finish"}\n\n'),
+        );
+        controller.close();
+      },
+    });
+    const response = new Response(body, {
+      status: 200,
+      headers: { "content-type": "text/event-stream" },
+    });
+    const fetchMock = vi.fn(async () => response);
+    const transport = new AssistantChatTransport({
+      resumable: { storage, resumeApi: "/api/resume" },
+      fetch: fetchMock as never,
+    });
+
+    const res = await wrappedFetchOf(transport)("https://example.com", {});
+
+    expect(res.status).toBe(200);
+    expect(res).not.toBe(response);
+    expect(res.body).not.toBeNull();
+    const reader = res.body!.getReader();
+    while (true) {
+      const { done } = await reader.read();
+      if (done) break;
+    }
+    expect(storage.getStreamId()).toBeNull();
+  });
+});

package/src/transport/AssistantChatTransport.ts

@@ -0,0 +1,204 @@
+import type {
+  AssistantRuntime,
+  ThreadListItemRuntime,
+} from "@assistant-ui/core";
+import {
+  DefaultChatTransport,
+  type HttpChatTransportInitOptions,
+  type UIMessage,
+} from "ai";
+import { toToolsJSONSchema } from "assistant-stream";
+import {
+  RESUMABLE_STREAM_ID_HEADER,
+  type AssistantChatResumableOptions,
+} from "./resumable";
+
+export type InitializableThreadListItem = Pick<
+  ThreadListItemRuntime,
+  "initialize"
+>;
+
+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";
+
+// 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.
+const NULL_BODY_STATUSES = new Set([101, 204, 205, 304]);
+
+export type AssistantChatTransportInitOptions<UI_MESSAGE extends UIMessage> =
+  HttpChatTransportInitOptions<UI_MESSAGE> & {
+    resumable?: AssistantChatResumableOptions;
+  };
+
+export class AssistantChatTransport<
+  UI_MESSAGE extends UIMessage,
+> extends DefaultChatTransport<UI_MESSAGE> {
+  private runtime: AssistantRuntime | undefined;
+  private getThreadListItem:
+    | (() => InitializableThreadListItem | undefined)
+    | undefined;
+  private readonly resumable: AssistantChatResumableOptions | undefined;
+
+  private readonly __internal_initOptions:
+    | AssistantChatTransportInitOptions<UI_MESSAGE>
+    | undefined;
+
+  /** Constructs an unwired copy with the same init options. */
+  public __internal_clone(): AssistantChatTransport<UI_MESSAGE> {
+    const Constructor = this.constructor as new (
+      initOptions?: AssistantChatTransportInitOptions<UI_MESSAGE>,
+    ) => AssistantChatTransport<UI_MESSAGE>;
+    return new Constructor(this.__internal_initOptions);
+  }
+
+  constructor(initOptions?: AssistantChatTransportInitOptions<UI_MESSAGE>) {
+    const { resumable, ...rest } = initOptions ?? {};
+    const userFetch = rest.fetch;
+    const userPrepareReconnect = rest.prepareReconnectToStreamRequest;
+
+    super({
+      ...rest,
+      ...(resumable && {
+        fetch: wrapFetchWithResumable(resumable, userFetch),
+        prepareReconnectToStreamRequest: wrapPrepareReconnect(
+          resumable,
+          userPrepareReconnect,
+        ),
+      }),
+      prepareSendMessagesRequest: async (options) => {
+        const threadId = options.id;
+        const context = this.runtime?.thread.getModelContext();
+        const threadListItem =
+          this.getThreadListItem?.() ?? this.runtime?.threads.mainItem;
+        const id = (await threadListItem?.initialize())?.remoteId ?? options.id;
+
+        const optionsEx = {
+          ...options,
+          id,
+          body: {
+            callSettings: context?.callSettings,
+            system: context?.system,
+            config: context?.config,
+            tools: toToolsJSONSchema(context?.tools ?? {}),
+            ...options?.body,
+          },
+        };
+        const preparedRequest =
+          await rest.prepareSendMessagesRequest?.(optionsEx);
+        const headers = resumable
+          ? new Headers(preparedRequest?.headers ?? options.headers)
+          : undefined;
+        headers?.set(RESUMABLE_THREAD_ID_HEADER, threadId);
+
+        return {
+          ...preparedRequest,
+          ...(headers && { headers }),
+          body: preparedRequest?.body ?? {
+            ...optionsEx.body,
+            id,
+            messages: options.messages,
+            trigger: options.trigger,
+            messageId: options.messageId,
+            metadata: options.requestMetadata,
+          },
+        };
+      },
+    });
+
+    this.resumable = resumable;
+    this.__internal_initOptions = initOptions;
+  }
+
+  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 headers = new Headers(init?.headers);
+    const threadId = headers.get(RESUMABLE_THREAD_ID_HEADER) ?? undefined;
+    headers.delete(RESUMABLE_THREAD_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);
+    if (!res.body || NULL_BODY_STATUSES.has(res.status)) 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)) {
+          if (!id || resumable.storage.getStreamId(threadId) === id) {
+            resumable.storage.clear(threadId);
+          }
+          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(options.id);
+    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 });
+    const headers = new Headers(userPrepared?.headers ?? options.headers);
+    headers.set(RESUMABLE_THREAD_ID_HEADER, options.id);
+    return {
+      ...userPrepared,
+      headers,
+      api: userPrepared?.api ?? api,
+    };
+  };
+}