@ably/ai-transport 0.1.0 → 0.3.0

package/src/vercel/react/contexts/chat-transport-provider.tsx

@@ -1,92 +1,100 @@
 /**
- * ChatTransportProvider: creates a ChatTransport from a ClientTransport and makes it
+ * ChatTransportProvider: creates a ChatTransport from a ClientSession and makes it
  * available to descendants via ChatTransportContext.
  *
- * Wraps children with TransportProvider (using UIMessageCodec) so the Ably channel
- * lifecycle is managed in one place. An inner component reads the ClientTransport
- * from NearestTransportContext and creates the ChatTransport once on first render
- * (via useRef).
+ * Wraps children with ClientSessionProvider (using UIMessageCodec). The
+ * surrounding `<AblyProvider>` supplies the Realtime client; the session
+ * resolves the channel from `channelName` itself. An inner component reads
+ * the ClientSession via useClientSession() and creates the ChatTransport
+ * via useMemo, keyed on the session and transport options.
  *
- * The ChatTransport is NOT closed on unmount — the underlying ClientTransport
- * lifecycle is managed by the wrapping TransportProvider. Auto-closing would break
- * React Strict Mode, and ChatTransport.close() delegates to ClientTransport.close()
- * which TransportProvider already calls.
+ * The ChatTransport is NOT closed on unmount — the underlying ClientSession
+ * lifecycle is managed by the wrapping ClientSessionProvider. Auto-closing would break
+ * React Strict Mode, and ChatTransport.close() delegates to ClientSession.close()
+ * which ClientSessionProvider already calls.
  *
  * Multiple ChatTransportProviders can be nested using distinct channelNames.
- * Each provider merges its transport into the parent registry, so descendants
- * can access all registered transports via useChatTransport({ channelName }).
+ * Each provider merges its session into the parent registry, so descendants
+ * can access all registered sessions via useChatTransport({ channelName }).
  */
 
 import type * as AI from 'ai';
 import { type PropsWithChildren, type ReactNode, useContext, useMemo } from 'react';
 
-import { createTransportHooks, type TransportProviderProps } from '../../../react/index.js';
-import { UIMessageCodec } from '../../codec/index.js';
-import { type ChatTransportOptions, DEFAULT_VERCEL_API } from '../../transport/index.js';
+import { type ClientSessionProviderProps, createSessionHooks } from '../../../react/index.js';
+import { UIMessageCodec, type VercelInput, type VercelOutput, type VercelProjection } from '../../codec/index.js';
+import type { ChatTransportOptions } from '../../transport/index.js';
 import { createChatTransport } from '../../transport/index.js';
 import type { ChatTransportSlot } from './chat-transport-context.js';
 import { ChatTransportContext } from './chat-transport-context.js';
 
-export const {
-  TransportProvider,
-  useAblyMessages,
-  useActiveTurns,
-  useClientTransport,
-  useCreateView,
-  useTree,
-  useView,
-} = createTransportHooks<AI.UIMessageChunk, AI.UIMessage>();
+export const { ClientSessionProvider, useAblyMessages, useClientSession, useCreateView, useTree, useView } =
+  createSessionHooks<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>();
 
-type CoreTransportProviderProps = Omit<TransportProviderProps<AI.UIMessageChunk, AI.UIMessage>, 'codec' | 'api'> &
-  Partial<Pick<TransportProviderProps<AI.UIMessageChunk, AI.UIMessage>, 'api'>>;
+type CoreClientSessionProviderProps = Omit<
+  ClientSessionProviderProps<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>,
+  'codec'
+>;
 
 /**
  * Props for {@link ChatTransportProvider}.
  *
- * All {@link TransportProviderProps} for Vercel types except `codec` (baked as UIMessageCodec),
- * plus `chatOptions` for customizing chat request construction.
+ * All {@link ClientSessionProviderProps} for Vercel types except `codec` (baked as UIMessageCodec),
+ * plus the transport-owned invocation POST options (`api` / `credentials` / `fetch`) and
+ * `chatOptions` for customizing chat request construction.
  */
-export interface ChatTransportProviderProps extends CoreTransportProviderProps {
+export interface ChatTransportProviderProps extends CoreClientSessionProviderProps {
+  /** Endpoint the chat transport POSTs the invocation to, to wake the agent. Default `/api/chat`. */
+  api?: string;
+  /** Fetch credentials mode for the invocation POST. */
+  credentials?: RequestCredentials;
+  /** Custom fetch implementation for the invocation POST. Defaults to `globalThis.fetch`. */
+  fetch?: typeof globalThis.fetch;
   /**
-   * Optional hooks for customizing chat request construction (e.g. prepareSendMessagesRequest).
+   * Optional transport options for customizing chat request construction (e.g. the `prepareSendMessagesRequest` hook).
    * Must be stable across renders — wrap in `useMemo` or define outside the component.
    * A new object reference triggers ChatTransport recreation.
+   * If this object also sets `api`/`credentials`/`fetch`, the dedicated top-level props of the same name take precedence.
    */
   chatOptions?: ChatTransportOptions;
 }
 
 const ChatTransportProviderInner = ({
   channelName,
-  chatOptions,
+  chatTransportOptions,
   children,
 }: {
   channelName: string;
-  chatOptions?: ChatTransportOptions;
+  chatTransportOptions: ChatTransportOptions;
   children: ReactNode;
 }) => {
-  const { transport, transportError } = useClientTransport();
+  const { session, sessionError } = useClientSession();
   const { providers: parentProviders } = useContext(ChatTransportContext);
-  const chatTransport = useMemo(() => createChatTransport(transport, chatOptions), [transport, chatOptions]);
+  const chatTransport = useMemo(
+    () => createChatTransport(session, chatTransportOptions),
+    [session, chatTransportOptions],
+  );
   const contextValue = useMemo(() => {
-    const slot: ChatTransportSlot = { transport, transportError, chatTransport };
+    const slot: ChatTransportSlot = { session, sessionError, chatTransport };
     return {
       nearest: slot,
       providers: { ...parentProviders, [channelName]: slot },
     };
-  }, [channelName, parentProviders, chatTransport, transport, transportError]);
+  }, [channelName, parentProviders, chatTransport, session, sessionError]);
 
   return <ChatTransportContext.Provider value={contextValue}>{children}</ChatTransportContext.Provider>;
 };
 
 /**
- * Provide a {@link ChatTransport} and its underlying {@link ClientTransport} to descendant components.
+ * Provide a {@link ChatTransport} and its underlying {@link ClientSession} to descendant components.
  *
- * Wraps children with Ably's `ChannelProvider` (via `TransportProvider`) using `channelName`,
- * creates a {@link ClientTransport} with UIMessageCodec, wraps it in a {@link ChatTransport},
+ * Wraps children with `ClientSessionProvider` using `channelName` (the Realtime
+ * client is read from the surrounding `<AblyProvider>`), creates a
+ * {@link ClientSession} with UIMessageCodec, wraps it in a {@link ChatTransport},
  * and registers the full slot in `ChatTransportContext` under `channelName`. Descendants call
- * {@link useChatTransport} with the same `channelName` to access both transports.
+ * {@link useChatTransport} with the same `channelName` to access both.
  *
- * `useClientTransport` is also available inside this provider's subtree.
+ * `useClientSession` is also available inside this provider's subtree.
  *
  * ```tsx
  * <ChatTransportProvider channelName="ai:demo">
@@ -94,29 +102,49 @@ const ChatTransportProviderInner = ({
  * </ChatTransportProvider>
  *
  * // Inside Chat:
- * const { chatTransport, transport } = useChatTransport();
- * const { transport } = useClientTransport(); // also available
+ * const { chatTransport, session } = useChatTransport();
+ * const { session } = useClientSession(); // also available
  * ```
- * @param props - Provider configuration including `channelName`, optional `chatOptions`, and all other transport options.
+ * @param props - Provider configuration including `channelName`, the invocation POST options (`api` / `credentials` / `fetch`), optional `chatOptions`, and all other session options.
+ * @param props.api - Endpoint the chat transport POSTs the invocation to. Default `/api/chat`.
+ * @param props.credentials - Fetch credentials mode for the invocation POST.
+ * @param props.fetch - Custom fetch implementation for the invocation POST.
  * @param props.chatOptions - Optional hooks for customizing chat request construction. Must be stable (memoized) — a new reference recreates the ChatTransport.
- * @param props.children - Descendant components that consume the transport via hooks.
- * @returns A React element wrapping children with ChannelProvider, TransportContext, and ChatTransportContext.
+ * @param props.children - Descendant components that consume the chat transport via hooks.
+ * @returns A React element wrapping children with ClientSessionContext and ChatTransportContext.
  */
 export const ChatTransportProvider = ({
+  api,
+  credentials,
+  fetch,
   chatOptions,
   children,
-  ...transportProps
-}: ChatTransportProviderProps & PropsWithChildren): ReactNode => (
-  <TransportProvider
-    {...transportProps}
-    api={transportProps.api ?? DEFAULT_VERCEL_API}
-    codec={UIMessageCodec}
-  >
-    <ChatTransportProviderInner
-      channelName={transportProps.channelName}
-      chatOptions={chatOptions}
+  ...sessionProps
+}: ChatTransportProviderProps & PropsWithChildren): ReactNode => {
+  // Fold the transport-owned POST options into a single ChatTransportOptions.
+  // Memoized so the ChatTransport isn't recreated each render — createChatTransport
+  // is keyed on this object's identity.
+  const chatTransportOptions = useMemo<ChatTransportOptions>(
+    () => ({
+      ...chatOptions,
+      ...(api !== undefined && { api }),
+      ...(credentials !== undefined && { credentials }),
+      ...(fetch !== undefined && { fetch }),
+    }),
+    [api, credentials, fetch, chatOptions],
+  );
+
+  return (
+    <ClientSessionProvider
+      {...sessionProps}
+      codec={UIMessageCodec}
     >
-      {children}
-    </ChatTransportProviderInner>
-  </TransportProvider>
-);
+      <ChatTransportProviderInner
+        channelName={sessionProps.channelName}
+        chatTransportOptions={chatTransportOptions}
+      >
+        {children}
+      </ChatTransportProviderInner>
+    </ClientSessionProvider>
+  );
+};

package/src/vercel/react/index.ts

@@ -1,12 +1,11 @@
-// Vercel-specific React hooks
+// Vercel-specific React entry point: providers, hooks, and their types
 export type { ChatTransport } from '../transport/chat-transport.js';
 export type { ChatTransportProviderProps } from './contexts/chat-transport-provider.js';
 export {
   ChatTransportProvider,
-  TransportProvider,
+  ClientSessionProvider,
   useAblyMessages,
-  useActiveTurns,
-  useClientTransport,
+  useClientSession,
   useCreateView,
   useTree,
   useView,
@@ -15,4 +14,3 @@ export type { ChatTransportHandle, UseChatTransportOptions } from './use-chat-tr
 export { useChatTransport } from './use-chat-transport.js';
 export type { UseMessageSyncOptions } from './use-message-sync.js';
 export { useMessageSync } from './use-message-sync.js';
-export { useStagedAddToolApprovalResponse } from './use-staged-add-tool-approval-response.js';

package/src/vercel/react/use-chat-transport.ts

@@ -1,55 +1,28 @@
 /**
- * useChatTransport: reads a ChatTransport and its underlying ClientTransport from
+ * useChatTransport: reads a ChatTransport and its underlying ClientSession from
  * the nearest ChatTransportProvider.
  *
- * The transport is created by ChatTransportProvider, which also wraps the subtree
- * with TransportProvider and Ably's ChannelProvider. This hook is a thin context
- * reader — it does not create or manage any transport state.
+ * The chat transport is created by ChatTransportProvider, which wraps the subtree
+ * with ClientSessionProvider. The Ably Realtime client is read from the
+ * surrounding `<AblyProvider>`. This hook is a thin context reader — it does
+ * not create or manage any session/transport state.
  *
  * Pass `channelName` to look up a specific provider by name. Omit to use the nearest
  * provider in the tree. Pass `skip: true` to defer (e.g. when auth is not yet resolved)
- * — returns stub transports whose properties throw with a descriptive error.
+ * — returns stubs whose properties throw with a descriptive error.
  */
 
 import * as Ably from 'ably';
 import type * as AI from 'ai';
 import { useContext } from 'react';
 
-import type { ClientTransport, Tree, View } from '../../core/transport/types.js';
+import type { ClientSession } from '../../core/transport/types.js';
 import { ErrorCode } from '../../errors.js';
+import { makeSkippedClientSession } from '../../react/internal/skipped-session.js';
+import type { VercelInput, VercelOutput, VercelProjection } from '../codec/index.js';
 import type { ChatTransport } from '../transport/index.js';
 import { ChatTransportContext } from './contexts/chat-transport-context.js';
 
-const SKIPPED_CLIENT_TRANSPORT: ClientTransport<AI.UIMessageChunk, AI.UIMessage> = {
-  get tree(): Tree<AI.UIMessage> {
-    throw new Ably.ErrorInfo('unable to access tree; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  get view(): View<AI.UIMessageChunk, AI.UIMessage> {
-    throw new Ably.ErrorInfo('unable to access view; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  createView: (): View<AI.UIMessageChunk, AI.UIMessage> => {
-    throw new Ably.ErrorInfo('unable to create view; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  cancel: () => {
-    throw new Ably.ErrorInfo('unable to cancel; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  stageEvents: () => {
-    throw new Ably.ErrorInfo('unable to stage events; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  stageMessage: () => {
-    throw new Ably.ErrorInfo('unable to stage message; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  waitForTurn: () => {
-    throw new Ably.ErrorInfo('unable to wait for turn; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  on: () => {
-    throw new Ably.ErrorInfo('unable to subscribe; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-  close: () => {
-    throw new Ably.ErrorInfo('unable to close; hook is skipped', ErrorCode.InvalidArgument, 400);
-  },
-};
-
 const SKIPPED_CHAT_TRANSPORT: ChatTransport = {
   sendMessages: (): never => {
     throw new Ably.ErrorInfo('unable to send messages; hook is skipped', ErrorCode.InvalidArgument, 400);
@@ -76,76 +49,81 @@ const SKIPPED_CHAT_TRANSPORT: ChatTransport = {
 export interface UseChatTransportOptions {
   /** Channel name to look up; omit to use the nearest {@link ChatTransportProvider}. */
   channelName?: string;
-  /** When `true`, return stub transports that throw on any access. */
+  /** When `true`, return stubs that throw on any access. */
   skip?: boolean;
 }
 
 /**
  * The value returned by {@link useChatTransport}.
- * Provides both the underlying {@link ClientTransport} and the {@link ChatTransport}
+ * Provides both the underlying {@link ClientSession} and the {@link ChatTransport}
  * adapter for Vercel's useChat hook.
  */
 export interface ChatTransportHandle {
   /**
-   * The underlying client transport, also available via {@link useClientTransport}.
-   * A throwing stub when `skip` is `true`, when no matching {@link TransportProvider}
-   * was found in the tree, or when transport construction failed. Check `transportError` before use.
+   * The underlying client session, also available via {@link useClientSession}.
+   * A throwing stub when `skip` is `true`, when no matching {@link ClientSessionProvider}
+   * was found in the tree, or when session construction failed. Check `sessionError` before use.
    */
-  transport: ClientTransport<AI.UIMessageChunk, AI.UIMessage>;
+  session: ClientSession<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>;
 
   /**
    * The chat transport adapter for use with Vercel's `useChat` hook.
    *
-   * A throwing stub when `skip` is `true`, when no matching
-   * {@link ChatTransportProvider} was found in the tree, or when the underlying
-   * {@link ClientTransport} construction failed. Check both `chatTransportError`
-   * and `transportError` before use.
+   * A throwing stub when `skip` is `true` or when no matching
+   * {@link ChatTransportProvider} was found in the tree. When a provider is found
+   * but the underlying {@link ClientSession} failed to construct, this is the real
+   * transport and `sessionError` is set instead. Check `chatTransportError` and
+   * `sessionError` before use.
    */
   chatTransport: ChatTransport;
 
   /**
-   * Set when no matching {@link TransportProvider} was found, when transport
+   * Set when no matching {@link ClientSessionProvider} was found, when session
    * construction failed, and `skip` is `false`.
-   * `undefined` when the transport resolved successfully or when `skip` is `true`.
+   * `undefined` when the session resolved successfully or when `skip` is `true`.
    */
-  transportError?: Ably.ErrorInfo | undefined;
+  sessionError?: Ably.ErrorInfo | undefined;
   /**
-   * Set when no matching {@link ChatTransportProvider} was found or when transport
-   * construction failed, and `skip` is `false`.
-   * `undefined` when the transport resolved successfully or when `skip` is `true`.
+   * Set only when no matching {@link ChatTransportProvider} was found and `skip` is
+   * `false`.
+   * `undefined` when the chat transport resolved successfully (even if session
+   * construction failed — see `sessionError`) or when `skip` is `true`.
    */
   chatTransportError?: Ably.ErrorInfo | undefined;
 }
 
 /**
- * Access a {@link ChatTransport} and {@link ClientTransport} from the nearest {@link ChatTransportProvider}.
+ * Access a {@link ChatTransport} and {@link ClientSession} from the nearest {@link ChatTransportProvider}.
  *
  * When `channelName` is omitted, the innermost `ChatTransportProvider` in the tree is used.
- * When `skip` is `true`, returns stub transports whose every property and method throws
+ * When `skip` is `true`, returns stubs whose every property and method throws
  * an {@link Ably.ErrorInfo} — safe to hold in state before conditions are ready.
  * When no provider is found, returns stubs with `chatTransportError` set instead of throwing.
- * @param props - Options for selecting the transport.
+ * @param props - Options for selecting the chat transport.
  * @param props.channelName - The channel name passed to the enclosing `ChatTransportProvider`. Omit to use the nearest.
  * @param props.skip - When `true`, return stubs that throw on any access instead of reading from context.
- * @returns The `ChatTransportHandle` containing both the chat transport adapter and the underlying client transport.
+ * @returns The `ChatTransportHandle` containing both the chat transport adapter and the underlying client session.
  */
 export const useChatTransport = ({ channelName, skip }: UseChatTransportOptions = {}): ChatTransportHandle => {
   const { nearest, providers } = useContext(ChatTransportContext);
 
   if (skip) {
-    return { transport: SKIPPED_CLIENT_TRANSPORT, chatTransport: SKIPPED_CHAT_TRANSPORT };
+    return {
+      session: makeSkippedClientSession<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>(),
+      chatTransport: SKIPPED_CHAT_TRANSPORT,
+    };
   }
 
   if (channelName !== undefined) {
     const slot = providers[channelName];
     if (slot) {
-      return { transport: slot.transport, chatTransport: slot.chatTransport, transportError: slot.transportError };
+      return { session: slot.session, chatTransport: slot.chatTransport, sessionError: slot.sessionError };
     }
     return {
-      transport: SKIPPED_CLIENT_TRANSPORT,
+      session: makeSkippedClientSession<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>(),
       chatTransport: SKIPPED_CHAT_TRANSPORT,
-      transportError: new Ably.ErrorInfo(
-        `unable to use client transport; no TransportProvider found for channelName "${channelName}"`,
+      sessionError: new Ably.ErrorInfo(
+        `unable to use client session; no ClientSessionProvider found for channelName "${channelName}"`,
         ErrorCode.BadRequest,
         400,
       ),
@@ -159,17 +137,17 @@ export const useChatTransport = ({ channelName, skip }: UseChatTransportOptions
 
   if (nearest) {
     return {
-      transport: nearest.transport,
+      session: nearest.session,
       chatTransport: nearest.chatTransport,
-      transportError: nearest.transportError,
+      sessionError: nearest.sessionError,
     };
   }
 
   return {
-    transport: SKIPPED_CLIENT_TRANSPORT,
+    session: makeSkippedClientSession<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>(),
     chatTransport: SKIPPED_CHAT_TRANSPORT,
-    transportError: new Ably.ErrorInfo(
-      'unable to use transport; no TransportProvider found in the tree',
+    sessionError: new Ably.ErrorInfo(
+      'unable to use session; no ClientSessionProvider found in the tree',
       ErrorCode.BadRequest,
       400,
     ),

package/src/vercel/react/use-message-sync.ts

@@ -1,72 +1,103 @@
 /**
- * useMessageSync: wires transport message lifecycle events into useChat's setMessages.
+ * useMessageSync: wire view updates into useChat's setMessages.
  *
- * Subscribes to the transport view's 'update' event and replaces messages state
- * with the view's authoritative message list.
+ * During active own-run streams, setMessages is gated to avoid an
+ * ID-mismatch in useChat's write(). When the stream ends, the gate
+ * opens and the view is synced into useChat's overlay.
  *
- * When a ChatTransport is provided (resolved from the nearest ChatTransportProvider),
- * setMessages calls are gated during active own-turn streams. This prevents the
- * push/replace ID mismatch in useChat's write() function. When the stream finishes,
- * the gate opens and an immediate sync fires to pick up any observer messages that
- * arrived during the stream.
- *
- * All dependencies are resolved from the nearest ChatTransportProvider via
- * useChatTransport(). Pass channelName to select a specific provider; omit to use
- * the nearest. Pass skip: true to pause all subscriptions.
- *
- * Returns the unsubscribe function in the useEffect cleanup so handlers
- * are removed on unmount or when dependencies change.
+ * The sync is a per-message merge, not a replace: when the overlay has
+ * resolved a client-side tool locally (via addToolResult) but the
+ * tree's echo hasn't landed yet, the overlay's resolution wins.
+ * Without that, the gate-open sync would race the AI SDK's post-stream
+ * sendAutomaticallyWhen check and could clobber the resolution before
+ * the continuation publishes.
  */
 
 import type * as AI from 'ai';
 import { useEffect, useState } from 'react';
 
+import { isToolPart, type ToolPart } from '../tool-part.js';
 import { useChatTransport } from './use-chat-transport.js';
 
 /** Options for {@link useMessageSync}. */
 export interface UseMessageSyncOptions {
   /**
-   * The `setMessages` updater function from `useChat()`. Required.
-   * Called with a function that replaces the previous message list with the
-   * transport's current authoritative message list.
+   * The `setMessages` updater function from `useChat()`. Called with an
+   * updater that returns the next overlay.
    */
   setMessages: (updater: (prev: AI.UIMessage[]) => AI.UIMessage[]) => void;
   /**
    * Channel name of the {@link ChatTransportProvider} to observe.
-   * Omit to use the nearest provider in the tree.
+   * Omit to use the nearest provider.
    */
   channelName?: string;
-  /**
-   * When `true`, skip all subscriptions and do nothing.
-   * Use when the hook's dependencies are not yet resolved (e.g. auth pending).
-   */
+  /** When `true`, skip all subscriptions. */
   skip?: boolean;
 }
 
+// ---------------------------------------------------------------------------
+// Tool-resolution merge
+// ---------------------------------------------------------------------------
+//
+// The merge matches tool parts by toolCallId (via the shared {@link isToolPart}
+// guard, which accepts both the codec's `dynamic-tool` shape and the AI SDK's
+// `tool-${name}` shape) and keeps the tree's `type` on the result so downstream
+// consumers narrowing on `dynamic-tool` keep working.
+
+const RESOLVED_TOOL_STATES = new Set(['output-available', 'output-error', 'approval-responded', 'output-denied']);
+
+const mergeAssistant = (tree: AI.UIMessage, overlay: AI.UIMessage): AI.UIMessage => {
+  const overlayByCallId = new Map<string, ToolPart>();
+  for (const part of overlay.parts) {
+    if (isToolPart(part)) overlayByCallId.set(part.toolCallId, part);
+  }
+  if (overlayByCallId.size === 0) return tree;
+
+  const parts = tree.parts.map((part) => {
+    if (!isToolPart(part)) return part;
+    if (RESOLVED_TOOL_STATES.has(part.state)) return part;
+    const overlayPart = overlayByCallId.get(part.toolCallId);
+    if (!overlayPart || !RESOLVED_TOOL_STATES.has(overlayPart.state)) return part;
+    // CAST: tool-${name} and dynamic-tool share the discriminated payload schema.
+    return { ...overlayPart, type: part.type } as AI.UIMessage['parts'][number];
+  });
+
+  const changed = parts.some((p, i) => p !== tree.parts[i]);
+  return changed ? { ...tree, parts } : tree;
+};
+
+const mergeMessages = (tree: AI.UIMessage[], overlay: AI.UIMessage[]): AI.UIMessage[] => {
+  if (overlay.length === 0) return tree;
+  const overlayById = new Map(overlay.map((m) => [m.id, m]));
+  return tree.map((treeMsg) => {
+    if (treeMsg.role !== 'assistant') return treeMsg;
+    const overlayMsg = overlayById.get(treeMsg.id);
+    return overlayMsg ? mergeAssistant(treeMsg, overlayMsg) : treeMsg;
+  });
+};
+
+// ---------------------------------------------------------------------------
+// Hook
+// ---------------------------------------------------------------------------
+
 /**
- * Wire transport message updates into `useChat()`'s `setMessages` updater.
- *
- * Resolves both the transport view and the streaming gate from the nearest
- * `ChatTransportProvider`. Pass `channelName` to target a specific provider.
- * Pass `skip: true` to pause all subscriptions.
+ * Subscribe to view updates and sync them into `useChat()`'s overlay.
  * @param options - Hook options.
- * @param options.setMessages - The `setMessages` function from `useChat()`. Required.
- * @param options.channelName - Channel name of the provider to observe; defaults to nearest.
+ * @param options.setMessages - The `setMessages` function from `useChat()`.
+ * @param options.channelName - Channel name of the provider to observe; defaults to the nearest.
  * @param options.skip - When `true`, skip all subscriptions.
  */
 export const useMessageSync = ({ setMessages, channelName, skip }: UseMessageSyncOptions): void => {
-  const { transport, chatTransport, chatTransportError } = useChatTransport({ channelName, skip });
+  const { session, chatTransport, chatTransportError } = useChatTransport({ channelName, skip });
 
-  // Only use resolved values when a provider was found and skip is false.
   const resolved = !skip && !chatTransportError;
-  const view = resolved ? transport.view : undefined;
+  const view = resolved ? session.view : undefined;
   const resolvedChatTransport = resolved ? chatTransport : undefined;
 
   const [gated, setGated] = useState(false);
 
-  // Subscribe to the ChatTransport's streaming state to gate setMessages.
-  // Reset gated to the new instance's current state so a stale `true`
-  // from a previous instance doesn't permanently suppress syncs.
+  // Subscribe to the ChatTransport's streaming state. Reset on transport
+  // change so a stale `true` doesn't permanently suppress syncs.
   useEffect(() => {
     if (!resolvedChatTransport) {
       setGated(false);
@@ -76,15 +107,20 @@ export const useMessageSync = ({ setMessages, channelName, skip }: UseMessageSyn
     return resolvedChatTransport.onStreamingChange(setGated);
   }, [resolvedChatTransport]);
 
-  // Subscribe to view updates and sync messages, unless gated.
+  // Subscribe to view updates and sync, unless gated.
   useEffect(() => {
     if (!view || gated) return;
 
     const sync = (): void => {
-      setMessages(() => view.flattenNodes().map((n) => n.message));
+      setMessages((overlay) =>
+        mergeMessages(
+          view.getMessages().map((m) => m.message),
+          overlay,
+        ),
+      );
     };
 
-    // Sync immediately when the effect runs (covers gate-open and initial mount).
+    // Sync immediately to cover gate-open and initial mount.
     sync();
 
     return view.on('update', sync);