@assistant-ui/ai-sdk 0.0.1

package/src/converters/toolOutputConversion.ts

@@ -0,0 +1,26 @@
+import type { JSONValue } from "ai";
+import type { ToolModelContentPart } from "assistant-stream";
+
+export const toAISDKContent = (parts: readonly ToolModelContentPart[]) => ({
+  type: "content" as const,
+  value: parts.map((part) => {
+    if (part.type === "text") {
+      return { type: "text" as const, text: part.text };
+    }
+    const mediaType =
+      typeof part.mediaType === "string"
+        ? part.mediaType
+        : "application/octet-stream";
+    return {
+      type: "file" as const,
+      data: { type: "data" as const, data: part.data },
+      mediaType,
+      ...(part.filename !== undefined && { filename: part.filename }),
+    };
+  }),
+});
+
+export const toAISDKDefaultOutput = (output: unknown) =>
+  typeof output === "string"
+    ? { type: "text" as const, value: output }
+    : { type: "json" as const, value: (output ?? null) as JSONValue };

package/src/hooks.ts

@@ -0,0 +1,23 @@
+"use client";
+
+import type { UseChatHelpers } from "@ai-sdk/react";
+import type { UIMessage } from "ai";
+import { aiSDKExtras } from "./aiSDKExtras";
+
+/**
+ * Read the last AI SDK chat error object from the runtime extras.
+ * `undefined` when there is no error or the current thread is not backed by
+ * the AI SDK runtime.
+ */
+export const useAISDKError = () => aiSDKExtras.use((e) => e.error, undefined);
+
+/**
+ * The underlying `useChat` helpers object, for advanced views — reach
+ * `resumeStream`, `clearError`, or the raw `sendMessage` without forking the
+ * runtime hook. `undefined` when the current thread is not backed by the AI
+ * SDK runtime.
+ */
+export const useAISDKChat = <UI_MESSAGE extends UIMessage = UIMessage>() =>
+  aiSDKExtras.use((e) => e.chat, undefined) as
+    | UseChatHelpers<UI_MESSAGE>
+    | undefined;

package/src/index.native.ts

@@ -0,0 +1,21 @@
+/// <reference types="@assistant-ui/core/react" />
+
+export { useAISDKRuntime } from "./runtime/useAISDKRuntime";
+export { useChatRuntime } from "./runtime/useChatRuntime";
+export type { UseChatRuntimeOptions } from "./runtime/useChatRuntime";
+export { AssistantChatTransport } from "./transport/AssistantChatTransport";
+export {
+  RESUMABLE_STREAM_ID_HEADER,
+  createResumableSessionStorage,
+} from "./transport/resumable";
+export type {
+  AssistantChatResumableOptions,
+  ResumableClientStorage,
+} from "./transport/resumable";
+export { frontendTools, type FrontendTools } from "./tools/frontendTools";
+export { injectQuoteContext } from "./model-context/injectQuoteContext";
+export { unstable_injectInteractableContext } from "./model-context/injectInteractableContext";
+export type { ThreadTokenUsage, TokenUsageExtractableMessage } from "./usage";
+export { getThreadMessageTokenUsage, useThreadTokenUsage } from "./usage";
+export { useAISDKChat, useAISDKError } from "./hooks";
+export { AISDKChat, type AISDKChatOptions } from "./runtime/AISDKChat";

package/src/index.ts

@@ -0,0 +1,29 @@
+/// <reference types="@assistant-ui/core/react" />
+
+export { useAISDKRuntime } from "./runtime/useAISDKRuntime";
+export { useChatRuntime } from "./runtime/useChatRuntime";
+export type { UseChatRuntimeOptions } from "./runtime/useChatRuntime";
+export { AssistantChatTransport } from "./transport/AssistantChatTransport";
+export {
+  RESUMABLE_STREAM_ID_HEADER,
+  createResumableSessionStorage,
+} from "./transport/resumable";
+export type {
+  AssistantChatResumableOptions,
+  ResumableClientStorage,
+} from "./transport/resumable";
+export { frontendTools, type FrontendTools } from "./tools/frontendTools";
+export { injectQuoteContext } from "./model-context/injectQuoteContext";
+export { unstable_injectInteractableContext } from "./model-context/injectInteractableContext";
+export type { ThreadTokenUsage, TokenUsageExtractableMessage } from "./usage";
+export { getThreadMessageTokenUsage, useThreadTokenUsage } from "./usage";
+export { useAISDKChat, useAISDKError } from "./hooks";
+export {
+  AISDKToolkit,
+  generativeTools,
+  type AISDKToolkitOptions,
+  type AISDKToolkitToolsOptions,
+  type GenerativeToolsOptions,
+} from "./tools/generativeTools";
+export { AISDKChat, type AISDKChatOptions } from "./runtime/AISDKChat";
+export { AISDKThreads, type AISDKThreadsOptions } from "./runtime/AISDKThreads";

package/src/model-context/injectInteractableContext.test.ts

@@ -0,0 +1,167 @@
+import type { UIMessage } from "ai";
+import { describe, expect, it } from "vitest";
+import { unstable_injectInteractableContext } from "./injectInteractableContext";
+import { injectQuoteContext } from "./injectQuoteContext";
+
+type Interactable = {
+  name: string;
+  id: string;
+  state: unknown;
+  partial?: boolean;
+};
+
+const userMsg = (
+  interactables?: Interactable[],
+  extraCustom?: Record<string, unknown>,
+): UIMessage =>
+  ({
+    id: "m1",
+    role: "user",
+    parts: [{ type: "text", text: "hello" }],
+    ...(interactables || extraCustom
+      ? {
+          metadata: {
+            custom: {
+              ...extraCustom,
+              ...(interactables ? { interactables } : {}),
+            },
+          },
+        }
+      : {}),
+  }) as UIMessage;
+
+const textOf = (part: unknown) => (part as { text?: string }).text;
+const isStateText = (part: unknown) =>
+  (part as { type?: string }).type === "text" &&
+  (textOf(part) ?? "").startsWith("[Current state of");
+
+describe("unstable_injectInteractableContext", () => {
+  it("returns the message unchanged when there is no interactable metadata", () => {
+    const input = [userMsg()];
+    const out = unstable_injectInteractableContext(input);
+    expect(out[0]).toBe(input[0]);
+  });
+
+  it("skips non-user messages", () => {
+    const input = [
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [{ type: "text", text: "hi" }],
+        metadata: {
+          custom: {
+            interactables: [{ name: "note", id: "n1", state: { v: 1 } }],
+          },
+        },
+      } as UIMessage,
+    ];
+    const out = unstable_injectInteractableContext(input);
+    expect(out[0]).toBe(input[0]);
+  });
+
+  it("prepends a text part using the default format (id included for update_* addressing)", () => {
+    const out = unstable_injectInteractableContext([
+      userMsg([{ name: "note", id: "n1", state: { title: "Hi" } }]),
+    ]);
+    expect(textOf(out[0]!.parts[0])).toBe(
+      '[Current state of "note" (id: "n1"): {"title":"Hi"}]\n\n',
+    );
+    expect(textOf(out[0]!.parts[1])).toBe("hello");
+  });
+
+  it("formats a partial snapshot as changed fields", () => {
+    const out = unstable_injectInteractableContext([
+      userMsg([
+        { name: "note", id: "n1", state: { title: "Hi" }, partial: true },
+      ]),
+    ]);
+    expect(textOf(out[0]!.parts[0])).toBe(
+      '[State of "note" (id: "n1") changed — updated fields: {"title":"Hi"}; fields not listed are unchanged]\n\n',
+    );
+  });
+
+  it("joins multiple interactables with a newline", () => {
+    const out = unstable_injectInteractableContext([
+      userMsg([
+        { name: "note", id: "n1", state: { v: 1 } },
+        { name: "board", id: "b1", state: { v: 2 } },
+      ]),
+    ]);
+    expect(textOf(out[0]!.parts[0])).toBe(
+      '[Current state of "note" (id: "n1"): {"v":1}]\n' +
+        '[Current state of "board" (id: "b1"): {"v":2}]\n\n',
+    );
+  });
+
+  it("respects a custom format function", () => {
+    const out = unstable_injectInteractableContext(
+      [userMsg([{ name: "note", id: "n1", state: { v: 1 } }])],
+      (item) => `X:${item.name}`,
+    );
+    expect(textOf(out[0]!.parts[0])).toBe("X:note\n\n");
+  });
+
+  it("leaves messages with an empty interactables array unchanged", () => {
+    const input = [userMsg([])];
+    expect(unstable_injectInteractableContext(input)[0]).toBe(input[0]);
+  });
+
+  it("does not throw when a user message has no parts", () => {
+    const input = [
+      {
+        id: "m1",
+        role: "user",
+        metadata: {
+          custom: {
+            interactables: [{ name: "note", id: "n1", state: { v: 1 } }],
+          },
+        },
+      } as unknown as UIMessage,
+    ];
+    const run = () => unstable_injectInteractableContext(input);
+    expect(run).not.toThrow();
+    expect(textOf(run()[0]!.parts[0])).toBe(
+      '[Current state of "note" (id: "n1"): {"v":1}]\n\n',
+    );
+  });
+
+  it("is idempotent on its own (does not double-inject)", () => {
+    const once = unstable_injectInteractableContext([
+      userMsg([{ name: "note", id: "n1", state: { v: 1 } }]),
+    ]);
+    const twice = unstable_injectInteractableContext(once);
+    expect(twice[0]).toBe(once[0]);
+    expect(twice[0]!.parts.filter(isStateText)).toHaveLength(1);
+  });
+
+  describe("composition with injectQuoteContext", () => {
+    const both = () =>
+      userMsg([{ name: "note", id: "n1", state: { v: 1 } }], {
+        quote: { text: "quoted" },
+      });
+
+    it("stacks both injections in a fixed order over source messages", () => {
+      // Apply interactable first, then quote — quote ends up outermost.
+      const out = injectQuoteContext(
+        unstable_injectInteractableContext([both()]),
+      );
+      const parts = out[0]!.parts;
+      expect(textOf(parts[0])).toBe("> quoted\n\n");
+      expect(isStateText(parts[1])).toBe(true);
+      expect(textOf(parts[2])).toBe("hello");
+    });
+
+    it("is NOT idempotent across re-application: run injectors once over un-injected messages", () => {
+      // Each injector's `alreadyInjected` guard only inspects parts[0]. Once the
+      // other injector sits at parts[0], the guard misses and re-prepends. So the
+      // composition must run once over source messages, never re-run over output.
+      const once = injectQuoteContext(
+        unstable_injectInteractableContext([both()]),
+      );
+      const twice = injectQuoteContext(
+        unstable_injectInteractableContext(once),
+      );
+      expect(twice[0]!.parts.filter(isStateText)).toHaveLength(2);
+    });
+  });
+});

package/src/model-context/injectInteractableContext.ts

@@ -0,0 +1,61 @@
+import type { UIMessage } from "ai";
+import {
+  unstable_formatInteractableSnapshot,
+  unstable_getInteractableSnapshots,
+  type Unstable_InteractableSnapshotEntry,
+} from "@assistant-ui/core";
+
+/**
+ * Injects interactable state snapshots into messages as model-visible text.
+ *
+ * Mirrors {@link injectQuoteContext}: reads the frozen snapshot stamped on a user
+ * message's `metadata.custom.interactables` (by the interactables scope at send
+ * time) and prepends a text part. Run this in your route handler before
+ * `convertToModelMessages`, which otherwise ignores `metadata.custom`.
+ *
+ * Wording is consumer-owned — pass `format` to control how each snapshot reads.
+ * A snapshot may originate from a user edit or an agent `update_*` call, so the
+ * default phrasing is neutral. Keep the instance id visible in custom wording:
+ * the model needs it to address the `update_*` tool's `id` parameter. A custom
+ * `format` must also handle entries with `partial: true`, whose `state` carries
+ * only the fields that changed since the model's last known state.
+ *
+ * @example
+ * ```ts
+ * import { convertToModelMessages, streamText } from "ai";
+ * import { unstable_injectInteractableContext } from "@assistant-ui/react-ai-sdk";
+ *
+ * export async function POST(req: Request) {
+ *   const { messages } = await req.json();
+ *   const result = streamText({
+ *     model: myModel,
+ *     messages: await convertToModelMessages(unstable_injectInteractableContext(messages)),
+ *   });
+ *   return result.toUIMessageStreamResponse();
+ * }
+ * ```
+ */
+export function unstable_injectInteractableContext(
+  messages: UIMessage[],
+  format: (
+    item: Unstable_InteractableSnapshotEntry,
+  ) => string = unstable_formatInteractableSnapshot,
+): UIMessage[] {
+  return messages.map((msg) => {
+    if (msg.role !== "user") return msg;
+
+    const items = unstable_getInteractableSnapshots(msg);
+    if (!items?.length) return msg;
+
+    const text = `${items.map(format).join("\n")}\n\n`;
+
+    const alreadyInjected =
+      msg.parts?.[0]?.type === "text" && msg.parts[0].text === text;
+    if (alreadyInjected) return msg;
+
+    return {
+      ...msg,
+      parts: [{ type: "text" as const, text }, ...(msg.parts ?? [])],
+    };
+  });
+}

package/src/model-context/injectQuoteContext.test.ts

@@ -0,0 +1,33 @@
+import type { UIMessage } from "ai";
+import { describe, expect, it } from "vitest";
+import { injectQuoteContext } from "./injectQuoteContext";
+
+const textOf = (part: unknown) => (part as { text?: string }).text;
+
+describe("injectQuoteContext", () => {
+  it("prepends the quote as a markdown blockquote", () => {
+    const out = injectQuoteContext([
+      {
+        id: "m1",
+        role: "user",
+        parts: [{ type: "text", text: "hello" }],
+        metadata: { custom: { quote: { text: "quoted" } } },
+      } as unknown as UIMessage,
+    ]);
+    expect(textOf(out[0]!.parts[0])).toBe("> quoted\n\n");
+    expect(textOf(out[0]!.parts[1])).toBe("hello");
+  });
+
+  it("does not throw when a user message has no parts", () => {
+    const input = [
+      {
+        id: "m1",
+        role: "user",
+        metadata: { custom: { quote: { text: "quoted" } } },
+      } as unknown as UIMessage,
+    ];
+    const run = () => injectQuoteContext(input);
+    expect(run).not.toThrow();
+    expect(textOf(run()[0]!.parts[0])).toBe("> quoted\n\n");
+  });
+});

package/src/model-context/injectQuoteContext.ts

@@ -0,0 +1,62 @@
+import type { UIMessage } from "ai";
+
+const getQuoteText = (metadata: unknown): string | undefined => {
+  if (!metadata || typeof metadata !== "object") return undefined;
+
+  const custom = (metadata as Record<string, unknown>).custom;
+  if (!custom || typeof custom !== "object") return undefined;
+
+  const quote = (custom as Record<string, unknown>).quote;
+  if (!quote || typeof quote !== "object") return undefined;
+
+  const text = (quote as Record<string, unknown>).text;
+  return typeof text === "string" ? text : undefined;
+};
+
+/**
+ * Injects quote context into messages as markdown blockquotes.
+ *
+ * Use this in your route handler before `convertToModelMessages` so the LLM
+ * sees the quoted text that the user is referring to.
+ *
+ * @example
+ * ```ts
+ * import { convertToModelMessages, streamText } from "ai";
+ * import { injectQuoteContext } from "@assistant-ui/react-ai-sdk";
+ *
+ * export async function POST(req: Request) {
+ *   const { messages } = await req.json();
+ *   const result = streamText({
+ *     model: myModel,
+ *     messages: await convertToModelMessages(injectQuoteContext(messages)),
+ *   });
+ *   return result.toUIMessageStreamResponse();
+ * }
+ * ```
+ */
+export function injectQuoteContext(messages: UIMessage[]): UIMessage[] {
+  return messages.map((msg) => {
+    if (msg.role !== "user") return msg;
+
+    const text = getQuoteText(msg.metadata);
+    if (!text) return msg;
+
+    const blockquote = text
+      .split(/\r?\n/)
+      .map((line) => `> ${line}`)
+      .join("\n");
+
+    const alreadyInjected =
+      msg.parts?.[0]?.type === "text" &&
+      msg.parts[0].text === `${blockquote}\n\n`;
+    if (alreadyInjected) return msg;
+
+    return {
+      ...msg,
+      parts: [
+        { type: "text" as const, text: `${blockquote}\n\n` },
+        ...(msg.parts ?? []),
+      ],
+    };
+  });
+}

package/src/runtime/AISDKChat.integration.test.tsx

@@ -0,0 +1,47 @@
+// @vitest-environment jsdom
+
+import { useEffect } from "react";
+import { render, waitFor } from "@testing-library/react";
+import { StrictMode } from "react";
+import { AuiConfig, AuiProvider, useAui } from "@assistant-ui/store";
+import { flushTapSync } from "@assistant-ui/tap";
+import { describe, expect, it } from "vitest";
+import { AISDKChat } from "./AISDKChat";
+import { createCancellableTransport } from "./__tests__/controlled-transport";
+
+describe("AISDKChat React integration", () => {
+  it("does not treat React provider unmount as client destruction", async () => {
+    const { transport, getCancelCount, close } = createCancellableTransport();
+    let started = false;
+    let isRunning = () => false;
+
+    const SendOnMount = () => {
+      const aui = useAui();
+      isRunning = () => aui.thread.getState().isRunning;
+      useEffect(() => {
+        if (started) return;
+        started = true;
+        flushTapSync(() => aui.composer.setText("keep streaming"));
+        flushTapSync(() => aui.composer.send());
+      }, [aui]);
+      return null;
+    };
+
+    const view = render(
+      <StrictMode>
+        <AuiProvider config={AuiConfig({ threads: AISDKChat({ transport }) })}>
+          <SendOnMount />
+        </AuiProvider>
+      </StrictMode>,
+    );
+
+    await waitFor(() => {
+      expect(isRunning()).toBe(true);
+    });
+
+    view.unmount();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(getCancelCount()).toBe(0);
+    close();
+  });
+});

package/src/runtime/AISDKChat.react-compiler.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, it, vi } from "vitest";
+import { flushTapSync } from "@assistant-ui/tap";
+import { AuiConfig, createAssistantClient } from "@assistant-ui/store/client";
+import { AISDKChat } from "./AISDKChat";
+import { createControlledTransport } from "./__tests__/controlled-transport";
+
+// Runs the streaming round trip against the compiled dist entry (aliased in
+// vitest.react-compiler.config.ts), so the react-compiler transform and the
+// react-shim import rewrite are exercised, not just the source hooks.
+describe("AISDKChat compiled output", () => {
+  it("streams a chat round trip through the compiled entry", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKChat({ transport }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+
+      flushTapSync(() => aui.composer.setText("hi"));
+      flushTapSync(() => aui.composer.send());
+
+      await vi.waitFor(() => {
+        expect(aui.thread.getState().isRunning).toBe(true);
+      });
+
+      emit(
+        { type: "start" },
+        { type: "text-start", id: "t1" },
+        { type: "text-delta", id: "t1", delta: "hello " },
+        { type: "text-delta", id: "t1", delta: "world" },
+        { type: "text-end", id: "t1" },
+        { type: "finish" },
+      );
+      close();
+
+      await vi.waitFor(() => {
+        const state = aui.thread.getState();
+        expect(state.isRunning).toBe(false);
+        expect(state.messages).toHaveLength(2);
+        expect(state.messages.at(-1)?.content).toContainEqual(
+          expect.objectContaining({ type: "text", text: "hello world" }),
+        );
+      });
+    } finally {
+      handle.destroy();
+    }
+  });
+});