@assistant-ui/ai-sdk 0.0.10 → 0.0.12

package/src/runtime/DynamicChatTransport.ts

@@ -0,0 +1,273 @@
+import type { AssistantRuntime } from "@assistant-ui/core";
+import type { UIMessage } from "@ai-sdk/react";
+import type { ChatTransport } from "ai";
+import {
+  AssistantChatTransport,
+  type InitializableThreadListItem,
+} from "../transport/AssistantChatTransport";
+import type { ResumableClientStorage } from "../transport/resumable";
+import { getResumableAdapter } from "./getResumableAdapter";
+
+const resumedStreamIdsByStorage = new WeakMap<
+  ResumableClientStorage,
+  Set<string>
+>();
+
+export const getResumedStreamIds = (
+  storage: ResumableClientStorage | undefined,
+) => {
+  if (!storage) return new Set<string>();
+  let resumedStreamIds = resumedStreamIdsByStorage.get(storage);
+  if (!resumedStreamIds) {
+    resumedStreamIds = new Set();
+    resumedStreamIdsByStorage.set(storage, resumedStreamIds);
+  }
+  return resumedStreamIds;
+};
+
+type ResumableStorageSubscription = {
+  listener: () => void;
+  threadId: string | undefined;
+  unsubscribe: (() => void) | undefined;
+};
+
+class DynamicResumableStorage implements ResumableClientStorage {
+  private readonly subscriptions = new Set<ResumableStorageSubscription>();
+  private storage: ResumableClientStorage | undefined;
+  private resumedStreamIds: Set<string>;
+  private hasPendingNotification = false;
+
+  constructor(storage: ResumableClientStorage | undefined) {
+    this.storage = storage;
+    this.resumedStreamIds = getResumedStreamIds(storage);
+  }
+
+  public setStorage(storage: ResumableClientStorage | undefined) {
+    if (this.storage === storage) return;
+    this.storage = storage;
+    const nextResumedStreamIds = getResumedStreamIds(storage);
+    for (const streamId of this.resumedStreamIds) {
+      nextResumedStreamIds.add(streamId);
+    }
+    this.resumedStreamIds = nextResumedStreamIds;
+    this.hasPendingNotification = true;
+  }
+
+  public flushChange() {
+    if (!this.hasPendingNotification) return;
+    this.hasPendingNotification = false;
+    for (const subscription of this.subscriptions) {
+      subscription.unsubscribe?.();
+      subscription.unsubscribe = this.storage?.subscribe?.(
+        subscription.listener,
+        subscription.threadId,
+      );
+      subscription.listener();
+    }
+  }
+
+  public getResumedStreamIds() {
+    return this.resumedStreamIds;
+  }
+
+  public getStreamId(threadId?: string) {
+    return this.storage?.getStreamId(threadId) ?? null;
+  }
+
+  public setStreamId(id: string, threadId?: string) {
+    this.storage?.setStreamId(id, threadId);
+  }
+
+  public clear(threadId?: string) {
+    this.storage?.clear(threadId);
+  }
+
+  public subscribe(listener: () => void, threadId?: string) {
+    const subscription: ResumableStorageSubscription = {
+      listener,
+      threadId,
+      unsubscribe: this.storage?.subscribe?.(listener, threadId),
+    };
+    this.subscriptions.add(subscription);
+    return () => {
+      this.subscriptions.delete(subscription);
+      subscription.unsubscribe?.();
+    };
+  }
+}
+
+type ThreadTransportContext<UI_MESSAGE extends UIMessage> = {
+  owner: object;
+  sourceTransport?: ChatTransport<UI_MESSAGE> | undefined;
+  transport?: ChatTransport<UI_MESSAGE> | undefined;
+  runtime?: AssistantRuntime | undefined;
+  getThreadListItem?:
+    | (() => InitializableThreadListItem | undefined)
+    | undefined;
+};
+
+type ThreadTransportBinding = {
+  runtime: AssistantRuntime;
+  getThreadListItem: () => InitializableThreadListItem | undefined;
+};
+
+export class DynamicChatTransport<
+  UI_MESSAGE extends UIMessage,
+> implements ChatTransport<UI_MESSAGE> {
+  private readonly threadContexts = new Map<
+    string,
+    ThreadTransportContext<UI_MESSAGE>
+  >();
+  private transport: ChatTransport<UI_MESSAGE>;
+  private readonly resumableStorage: DynamicResumableStorage;
+
+  constructor(transport: ChatTransport<UI_MESSAGE>) {
+    this.transport = transport;
+    this.resumableStorage = new DynamicResumableStorage(
+      getResumableAdapter(transport)?.storage,
+    );
+  }
+
+  public readonly sendMessages: ChatTransport<UI_MESSAGE>["sendMessages"] = (
+    options,
+  ) => this.getTransport(options.chatId).sendMessages(options);
+
+  public readonly reconnectToStream: ChatTransport<UI_MESSAGE>["reconnectToStream"] =
+    (options) => this.getTransport(options.chatId).reconnectToStream(options);
+
+  public readonly getCurrentTransport = (chatId: string) =>
+    this.getThreadTransport(this.threadContexts.get(chatId)) ?? this.transport;
+
+  public readonly getCurrentResumableStorage = () => this.resumableStorage;
+
+  public readonly getResumedStreamIds = () =>
+    this.resumableStorage.getResumedStreamIds();
+
+  public createThreadProxy(
+    owner: object,
+    getBinding: () => ThreadTransportBinding,
+  ): ChatTransport<UI_MESSAGE> {
+    return {
+      sendMessages: (options) =>
+        this.getOrCreateBoundTransport(
+          options.chatId,
+          owner,
+          getBinding,
+        ).sendMessages(options),
+      reconnectToStream: (options) =>
+        this.getOrCreateBoundTransport(
+          options.chatId,
+          owner,
+          getBinding,
+        ).reconnectToStream(options),
+    };
+  }
+
+  public setTransport(transport: ChatTransport<UI_MESSAGE>) {
+    if (this.transport === transport) return;
+    this.transport = transport;
+    this.resumableStorage.setStorage(getResumableAdapter(transport)?.storage);
+  }
+
+  public flushTransportChange() {
+    this.resumableStorage.flushChange();
+  }
+
+  public registerThread(chatId: string, owner: object) {
+    const existing = this.threadContexts.get(chatId);
+    if (existing?.owner === owner) return;
+    this.threadContexts.set(chatId, {
+      owner,
+    });
+  }
+
+  public setThreadContext(
+    chatId: string,
+    owner: object,
+    runtime: AssistantRuntime,
+    getThreadListItem: () => InitializableThreadListItem | undefined,
+  ) {
+    this.registerThread(chatId, owner);
+    const context = this.threadContexts.get(chatId)!;
+    context.runtime = runtime;
+    context.getThreadListItem = getThreadListItem;
+    if (context.transport === undefined) {
+      this.getThreadTransport(context);
+      return;
+    }
+    if (
+      context.sourceTransport === this.transport &&
+      context.transport !== undefined
+    ) {
+      this.wireTransport(context, context.transport);
+    }
+  }
+
+  public unregisterThread(chatId: string, owner: object) {
+    if (this.threadContexts.get(chatId)?.owner === owner) {
+      this.threadContexts.delete(chatId);
+    }
+  }
+
+  private getTransport(chatId: string) {
+    const context = this.threadContexts.get(chatId);
+    if (!context) {
+      throw new Error(
+        `DynamicChatTransport has no registered context for chat "${chatId}"`,
+      );
+    }
+    return this.getThreadTransport(context)!;
+  }
+
+  private getBoundTransport(
+    chatId: string,
+    owner: object,
+    binding: ThreadTransportBinding,
+  ) {
+    this.registerThread(chatId, owner);
+    const context = this.threadContexts.get(chatId)!;
+    context.runtime = binding.runtime;
+    context.getThreadListItem = binding.getThreadListItem;
+    return this.getThreadTransport(context)!;
+  }
+
+  private getOrCreateBoundTransport(
+    chatId: string,
+    owner: object,
+    getBinding: () => ThreadTransportBinding,
+  ) {
+    const context = this.threadContexts.get(chatId);
+    return context?.owner === owner
+      ? this.getThreadTransport(context)!
+      : this.getBoundTransport(chatId, owner, getBinding());
+  }
+
+  private getThreadTransport(
+    context: ThreadTransportContext<UI_MESSAGE> | undefined,
+  ) {
+    if (!context) return undefined;
+    if (context.sourceTransport !== this.transport) {
+      context.transport = this.createThreadTransport(this.transport);
+      context.sourceTransport = this.transport;
+    }
+    this.wireTransport(context, context.transport!);
+    return context.transport!;
+  }
+
+  private createThreadTransport(transport: ChatTransport<UI_MESSAGE>) {
+    return transport instanceof AssistantChatTransport
+      ? transport.__internal_clone()
+      : transport;
+  }
+
+  private wireTransport(
+    context: ThreadTransportContext<UI_MESSAGE>,
+    transport: ChatTransport<UI_MESSAGE>,
+  ) {
+    if (!(transport instanceof AssistantChatTransport)) return;
+    if (context.runtime) transport.setRuntime(context.runtime);
+    if (context.getThreadListItem) {
+      transport.__internal_setGetThreadListItem(context.getThreadListItem);
+    }
+  }
+}

package/src/runtime/__tests__/controlled-transport.ts

@@ -40,6 +40,9 @@ export const createCancellableTransport = () => {
   return {
     transport,
     getCancelCount: () => cancelCount,
+    emit: (...chunks: UIMessageChunk[]) => {
+      for (const chunk of chunks) controller.enqueue(chunk);
+    },
     close: () => controller.close(),
   };
 };

package/src/runtime/getResumableAdapter.ts

@@ -0,0 +1,16 @@
+import type { UIMessage } from "@ai-sdk/react";
+import type { ChatTransport } from "ai";
+import { AssistantChatTransport } from "../transport/AssistantChatTransport";
+import type { AssistantChatResumableOptions } from "../transport/resumable";
+
+export 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;
+};

package/src/runtime/toolHistoryCodec.test.ts

@@ -0,0 +1,161 @@
+import type {
+  RespondToToolApprovalOptions,
+  ThreadMessage,
+  Unstable_ToolInteractionLog,
+} from "@assistant-ui/core";
+import { describe, expect, it } from "vitest";
+import {
+  addToolData,
+  collectToolApprovalResponses,
+  collectToolArtifacts,
+  collectToolInteractions,
+  restoreToolData,
+} from "./toolHistoryCodec";
+
+describe("tool history codec", () => {
+  it("round trips artifacts, interactions, and approval responses", () => {
+    const log: Unstable_ToolInteractionLog = {
+      entries: [{ type: "action", occurredAt: 1, payload: { refresh: true } }],
+    };
+    const response: RespondToToolApprovalOptions = {
+      approvalId: "approval-1",
+      approved: true,
+      answers: { scope: { optionIds: ["tests"] } },
+      reason: "Approved",
+    };
+    const threadMessage: ThreadMessage = {
+      id: "assistant-1",
+      role: "assistant",
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "call-1",
+          toolName: "weather",
+          args: {},
+          argsText: "{}",
+          result: undefined,
+          isError: false,
+          approval: { id: "approval-1" },
+        },
+      ],
+      createdAt: new Date(),
+      status: { type: "complete", reason: "stop" },
+      metadata: {
+        unstable_state: null,
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    };
+    const innerMessage = {
+      id: "inner-1",
+      parts: [{ toolCallId: "call-1", approval: { id: "approval-1" } }],
+      metadata: { custom: "kept" },
+    };
+    const artifacts = new Map<string, unknown>([
+      ["call-1", { preview: "sunny" }],
+      ["unrelated", "ignored"],
+    ]);
+    const interactions = new Map<string, Unstable_ToolInteractionLog>([
+      ["call-1", log],
+      ["unrelated", log],
+    ]);
+    const responses = new Map<string, RespondToToolApprovalOptions>([
+      ["approval-1", response],
+      ["unrelated", response],
+    ]);
+
+    const encoded = addToolData(
+      innerMessage,
+      collectToolArtifacts(threadMessage, artifacts),
+      collectToolInteractions(threadMessage, interactions),
+      collectToolApprovalResponses(threadMessage, responses),
+    );
+
+    expect(encoded).toEqual({
+      ...innerMessage,
+      metadata: {
+        custom: "kept",
+        __aui_toolArtifacts: { "call-1": { preview: "sunny" } },
+        __aui_toolInteractions: { "call-1": log },
+        __aui_toolApprovalResponses: {
+          "approval-1": {
+            approved: true,
+            answers: response.answers,
+            reason: "Approved",
+          },
+        },
+      },
+    });
+    expect(innerMessage.metadata).toEqual({ custom: "kept" });
+
+    const restoredArtifacts = new Map<string, unknown>();
+    const restoredInteractions = new Map<string, Unstable_ToolInteractionLog>();
+    const restoredResponses = new Map<string, RespondToToolApprovalOptions>();
+    expect(
+      restoreToolData(
+        encoded,
+        restoredArtifacts,
+        restoredInteractions,
+        restoredResponses,
+      ),
+    ).toEqual(innerMessage);
+    expect(restoredArtifacts).toEqual(
+      new Map([["call-1", { preview: "sunny" }]]),
+    );
+    expect(restoredInteractions).toEqual(new Map([["call-1", log]]));
+    expect(restoredResponses).toEqual(new Map([["approval-1", response]]));
+  });
+
+  it("leaves messages without tool metadata unchanged", () => {
+    const message = {
+      parts: [{ toolCallId: "call-1" }],
+      metadata: { custom: 1 },
+    };
+    expect(addToolData(message, undefined, undefined, undefined)).toBe(message);
+    expect(restoreToolData(message, new Map(), new Map(), new Map())).toBe(
+      message,
+    );
+    expect(
+      restoreToolData({ parts: [] }, new Map(), new Map(), new Map()),
+    ).toEqual({ parts: [] });
+  });
+
+  it("ignores non-object messages and malformed metadata values", () => {
+    expect(
+      addToolData(null, { "call-1": "artifact" }, undefined, undefined),
+    ).toBeNull();
+    const invalidParts = { parts: "invalid" };
+    expect(
+      addToolData(invalidParts, { "call-1": "artifact" }, undefined, undefined),
+    ).toBe(invalidParts);
+    expect(restoreToolData(42, new Map(), new Map(), new Map())).toBe(42);
+    const invalidMetadata = { parts: [], metadata: "invalid" };
+    expect(
+      restoreToolData(invalidMetadata, new Map(), new Map(), new Map()),
+    ).toBe(invalidMetadata);
+
+    const artifacts = new Map<string, unknown>();
+    const interactions = new Map<string, Unstable_ToolInteractionLog>();
+    const responses = new Map<string, RespondToToolApprovalOptions>();
+    expect(
+      restoreToolData(
+        {
+          metadata: {
+            custom: "kept",
+            __aui_toolArtifacts: "invalid",
+            __aui_toolInteractions: 1,
+            __aui_toolApprovalResponses: [],
+          },
+        },
+        artifacts,
+        interactions,
+        responses,
+      ),
+    ).toEqual({ metadata: { custom: "kept" } });
+    expect(artifacts.size).toBe(0);
+    expect(interactions.size).toBe(0);
+    expect(responses.size).toBe(0);
+  });
+});

package/src/runtime/toolHistoryCodec.ts

@@ -0,0 +1,207 @@
+import type {
+  RespondToToolApprovalOptions,
+  ThreadMessage,
+  Unstable_ToolInteractionLog,
+} from "@assistant-ui/core";
+import { isRecord, readToolInteractionLog } from "@assistant-ui/core/internal";
+import { normalizeToolApprovalAnswers } from "../converters/toolApprovalAnswers";
+
+const TOOL_ARTIFACTS_METADATA_KEY = "__aui_toolArtifacts";
+const TOOL_INTERACTIONS_METADATA_KEY = "__aui_toolInteractions";
+const TOOL_APPROVAL_RESPONSES_METADATA_KEY = "__aui_toolApprovalResponses";
+
+export type StoredToolApprovalResponse = Omit<
+  RespondToToolApprovalOptions,
+  "approvalId"
+>;
+
+export const collectToolArtifacts = (
+  message: ThreadMessage,
+  toolArtifacts: ReadonlyMap<string, unknown> | undefined,
+) => {
+  if (!toolArtifacts) return undefined;
+  const entries = message.content.flatMap((part) => {
+    if (part.type !== "tool-call") return [];
+    const artifact = toolArtifacts.get(part.toolCallId);
+    return artifact === undefined ? [] : [[part.toolCallId, artifact] as const];
+  });
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+};
+
+export const collectToolInteractions = (
+  message: ThreadMessage,
+  toolInteractions:
+    | ReadonlyMap<string, Unstable_ToolInteractionLog>
+    | undefined,
+) => {
+  if (!toolInteractions) return undefined;
+  const entries = message.content.flatMap((part) => {
+    if (part.type !== "tool-call") return [];
+    const interactions = toolInteractions.get(part.toolCallId);
+    return interactions === undefined
+      ? []
+      : [[part.toolCallId, interactions] as const];
+  });
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+};
+
+export const collectToolApprovalResponses = (
+  message: ThreadMessage,
+  toolApprovalResponses:
+    | ReadonlyMap<string, RespondToToolApprovalOptions>
+    | undefined,
+) => {
+  if (!toolApprovalResponses) return undefined;
+  const entries = message.content.flatMap((part) => {
+    if (part.type !== "tool-call" || !part.approval) return [];
+    const response = toolApprovalResponses.get(part.approval.id);
+    if (!response) return [];
+    return [
+      [
+        part.approval.id,
+        {
+          approved: response.approved,
+          ...(response.optionId != null && { optionId: response.optionId }),
+          ...(response.text != null && { text: response.text }),
+          ...(response.answers != null && { answers: response.answers }),
+          ...(response.reason != null && { reason: response.reason }),
+        },
+      ] as const,
+    ];
+  });
+  return entries.length > 0 ? Object.fromEntries(entries) : undefined;
+};
+
+export const addToolData = <TMessage>(
+  message: TMessage,
+  toolArtifacts: Record<string, unknown> | undefined,
+  toolInteractions: Record<string, Unstable_ToolInteractionLog> | undefined,
+  toolApprovalResponses: Record<string, StoredToolApprovalResponse> | undefined,
+): TMessage => {
+  if (
+    (!toolArtifacts && !toolInteractions && !toolApprovalResponses) ||
+    !isRecord(message) ||
+    !Array.isArray(message.parts)
+  )
+    return message;
+  const toolCallIds = message.parts.flatMap((part) => {
+    if (!isRecord(part) || typeof part.toolCallId !== "string") return [];
+    return [part.toolCallId];
+  });
+  const artifacts = toolArtifacts
+    ? Object.fromEntries(
+        toolCallIds.flatMap((toolCallId) =>
+          Object.hasOwn(toolArtifacts, toolCallId)
+            ? [[toolCallId, toolArtifacts[toolCallId]] as const]
+            : [],
+        ),
+      )
+    : undefined;
+  const interactions = toolInteractions
+    ? Object.fromEntries(
+        toolCallIds.flatMap((toolCallId) =>
+          Object.hasOwn(toolInteractions, toolCallId)
+            ? [[toolCallId, toolInteractions[toolCallId]] as const]
+            : [],
+        ),
+      )
+    : undefined;
+  const approvalIds = message.parts.flatMap((part) => {
+    if (!isRecord(part) || !isRecord(part.approval)) return [];
+    const approvalId = part.approval.id;
+    return typeof approvalId === "string" ? [approvalId] : [];
+  });
+  const approvalResponses = toolApprovalResponses
+    ? Object.fromEntries(
+        approvalIds.flatMap((approvalId) =>
+          Object.hasOwn(toolApprovalResponses, approvalId)
+            ? [[approvalId, toolApprovalResponses[approvalId]] as const]
+            : [],
+        ),
+      )
+    : undefined;
+  const hasArtifacts = !!artifacts && Object.keys(artifacts).length > 0;
+  const hasInteractions =
+    !!interactions && Object.keys(interactions).length > 0;
+  const hasApprovalResponses =
+    !!approvalResponses && Object.keys(approvalResponses).length > 0;
+  if (!hasArtifacts && !hasInteractions && !hasApprovalResponses)
+    return message;
+  const metadata = isRecord(message.metadata) ? message.metadata : {};
+  return {
+    ...message,
+    metadata: {
+      ...metadata,
+      ...(hasArtifacts && { [TOOL_ARTIFACTS_METADATA_KEY]: artifacts }),
+      ...(hasInteractions && {
+        [TOOL_INTERACTIONS_METADATA_KEY]: interactions,
+      }),
+      ...(hasApprovalResponses && {
+        [TOOL_APPROVAL_RESPONSES_METADATA_KEY]: approvalResponses,
+      }),
+    },
+  } as TMessage;
+};
+
+export const restoreToolData = <TMessage>(
+  message: TMessage,
+  toolArtifacts: Map<string, unknown> | undefined,
+  toolInteractions: Map<string, Unstable_ToolInteractionLog> | undefined,
+  toolApprovalResponses: Map<string, RespondToToolApprovalOptions> | undefined,
+): TMessage => {
+  if (!isRecord(message) || !isRecord(message.metadata)) return message;
+  const metadata = message.metadata;
+  const hasArtifacts = Object.hasOwn(metadata, TOOL_ARTIFACTS_METADATA_KEY);
+  const hasInteractions = Object.hasOwn(
+    metadata,
+    TOOL_INTERACTIONS_METADATA_KEY,
+  );
+  const hasApprovalResponses = Object.hasOwn(
+    metadata,
+    TOOL_APPROVAL_RESPONSES_METADATA_KEY,
+  );
+  if (!hasArtifacts && !hasInteractions && !hasApprovalResponses)
+    return message;
+  const artifacts = metadata[TOOL_ARTIFACTS_METADATA_KEY];
+  if (toolArtifacts && isRecord(artifacts)) {
+    for (const [toolCallId, artifact] of Object.entries(artifacts)) {
+      toolArtifacts.set(toolCallId, artifact);
+    }
+  }
+  const interactions = metadata[TOOL_INTERACTIONS_METADATA_KEY];
+  if (toolInteractions && isRecord(interactions)) {
+    for (const [toolCallId, value] of Object.entries(interactions)) {
+      const log = readToolInteractionLog(value);
+      if (log) toolInteractions.set(toolCallId, log);
+    }
+  }
+  const approvalResponses = metadata[TOOL_APPROVAL_RESPONSES_METADATA_KEY];
+  if (toolApprovalResponses && isRecord(approvalResponses)) {
+    for (const [approvalId, value] of Object.entries(approvalResponses)) {
+      if (!isRecord(value) || typeof value.approved !== "boolean") continue;
+      const answers = normalizeToolApprovalAnswers(value.answers);
+      toolApprovalResponses.set(approvalId, {
+        approvalId,
+        approved: value.approved,
+        ...(typeof value.optionId === "string" && {
+          optionId: value.optionId,
+        }),
+        ...(typeof value.text === "string" && { text: value.text }),
+        ...(answers && Object.keys(answers).length > 0 && { answers }),
+        ...(typeof value.reason === "string" && { reason: value.reason }),
+      });
+    }
+  }
+  const {
+    [TOOL_ARTIFACTS_METADATA_KEY]: _,
+    [TOOL_INTERACTIONS_METADATA_KEY]: __,
+    [TOOL_APPROVAL_RESPONSES_METADATA_KEY]: ___,
+    ...restMetadata
+  } = metadata;
+  const { metadata: _metadata, ...restMessage } = message;
+  return (
+    Object.keys(restMetadata).length === 0
+      ? restMessage
+      : { ...restMessage, metadata: restMetadata }
+  ) as TMessage;
+};

package/src/runtime/useAISDKRuntime.approval.test.tsx

@@ -203,6 +203,18 @@ describe("useAISDKRuntime tool approvals", () => {
     });
   });
 
+  it("finds the pending approval past malformed parts", async () => {
+    const onRespondToToolApproval = vi.fn(async () => {});
+    const { respond, messages } = setupPendingApproval(onRespondToToolApproval);
+    (messages[0]!.parts as unknown[]).unshift(null, { text: "no type" });
+
+    await act(async () => {
+      await respond({ approvalId: "approval-1", approved: true });
+    });
+
+    expect(onRespondToToolApproval).toHaveBeenCalledOnce();
+  });
+
   it("does not store a request the handler hands back through the AI SDK", async () => {
     const { respond, addToolApprovalResponse } = setupPendingApproval(
       (_response, { respondViaAISDK }) => respondViaAISDK(),