@ably/ai-transport 0.1.0 → 0.2.0

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

@@ -1,21 +1,16 @@
 /**
- * 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';
@@ -26,47 +21,88 @@ 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 Vercel codec normalises every tool part to `dynamic-tool`, but the
+// AI SDK emits `tool-${name}` for statically-declared tools. Both shapes
+// share `toolCallId` + `state`; the merge matches by toolCallId and keeps
+// the tree's `type` on the result so downstream consumers narrowing on
+// `dynamic-tool` keep working.
+
+type ToolPart = AI.DynamicToolUIPart | AI.ToolUIPart;
+
+const RESOLVED_TOOL_STATES = new Set(['output-available', 'output-error', 'approval-responded', 'output-denied']);
+
+const isToolPart = (part: AI.UIMessage['parts'][number]): part is ToolPart =>
+  (part.type === 'dynamic-tool' || part.type.startsWith('tool-')) && 'toolCallId' in part && 'state' in part;
+
+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 +112,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);

package/src/vercel/run-end-reason.ts

@@ -0,0 +1,78 @@
+import type * as AI from 'ai';
+
+import type { RunEndReason, StreamResult } from '../core/transport/types.js';
+
+/**
+ * Derive the outcome for a Vercel `streamText` response that was piped through
+ * `Run.pipe`: either a terminal {@link RunEndReason} the caller passes to
+ * `Run.end`, or the sentinel `'suspend'` telling the caller to call
+ * `Run.suspend` instead. Preserves transport-level outcomes (`'cancelled'`,
+ * `'error'`) from the pipe result; when the pipe completed naturally, awaits
+ * Vercel's `finishReason` and returns `'suspend'` for `'tool-calls'` (the LLM
+ * requested tools the SDK did not auto-execute, so the run should suspend
+ * rather than end), or `'complete'` otherwise.
+ *
+ * Tolerates `finishReason` rejection. Vercel AI SDK v6 rejects
+ * `streamText().finishReason` with the abort signal's reason when the stream
+ * is aborted before any step completes, and rejects with
+ * `NoOutputGeneratedError` when the model produced nothing at all. Without
+ * this guard the rejection would bubble out of the route handler's `after()`
+ * block, skip the developer's `Run.end(...)` call, and leave the run with no
+ * `ai-run-end` event on the channel — so observers' UIs stay stuck on
+ * `streaming` indefinitely.
+ *
+ * Saves callers from interpreting Vercel domain semantics inline at the end
+ * of every route handler.
+ * @param pipeResult - The result returned by `Run.pipe`.
+ * @param finishReason - The `finishReason` promise from a `streamText` result.
+ * @returns `'suspend'` when the run should suspend awaiting tool input, or the
+ *   {@link RunEndReason} to pass to `Run.end` otherwise.
+ */
+export const vercelRunOutcome = async (
+  pipeResult: StreamResult,
+  finishReason: PromiseLike<AI.FinishReason>,
+): Promise<RunEndReason | 'suspend'> => {
+  if (pipeResult.reason !== 'complete') {
+    // Vercel's `result.finishReason` getter creates the underlying Promise
+    // eagerly, before the caller hands it to us. When `streamText` is
+    // aborted before any step completes, Vercel rejects that Promise with
+    // the abort signal's reason — typically a DOMException whose
+    // `.message` is a read-only getter. Returning early without ever
+    // attaching a handler lets Node report it as an unhandled rejection;
+    // Next.js' dev bundler then tries to mutate `.message` for logging
+    // and crashes with a confusing TypeError. Attach a silent handler so
+    // the rejection is observed and discarded — the transport-level
+    // `pipeResult.reason` is already what we return.
+    Promise.resolve(finishReason).catch(() => {
+      /* intentionally discarded; reason already known from pipeResult */
+    });
+    return pipeResult.reason;
+  }
+  try {
+    const finish = await finishReason;
+    return finish === 'tool-calls' ? 'suspend' : 'complete';
+  } catch (error) {
+    // Abort-shaped rejections are surfaced from streamText when the run was
+    // cancelled before any step finished — treat the run as cancelled so the
+    // observable lifecycle matches the cancel that triggered it. Everything
+    // else is a real error (e.g. NoOutputGeneratedError, network blow-ups);
+    // surface it as such so the developer sees the failure rather than a
+    // silent cancel.
+    return _isAbortLikeError(error) ? 'cancelled' : 'error';
+  }
+};
+
+/**
+ * Heuristic for "this error came from an AbortSignal aborting".
+ * Covers `DOMException` aborts (browser / Node 20+ `streamText`),
+ * plain `Error` objects whose `name` is `'AbortError'`, and anything
+ * else carrying that conventional name. Avoids importing
+ * `@ai-sdk/provider-utils` just for `isAbortError`.
+ * @param error - The error to test.
+ * @returns `true` if the error looks like an abort.
+ */
+const _isAbortLikeError = (error: unknown): boolean => {
+  if (typeof error !== 'object' || error === null) return false;
+  const name = (error as { name?: unknown }).name;
+  return name === 'AbortError';
+};