@ably/ai-transport 0.0.1 → 0.2.0

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

@@ -1,34 +1,133 @@
 /**
- * useMessageSync: wires transport message lifecycle events into useChat's setMessages.
+ * useMessageSync: wire view updates into useChat's setMessages.
  *
- * Subscribes to the transport's 'message' event and replaces messages state
- * with the transport's authoritative message list. Events fire immediately
- * on every store update (including during active streaming), so this hook
- * keeps React state in sync in real time.
+ * 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.
  *
- * 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 } from 'react';
+import { useEffect, useState } from 'react';
 
-import type { ClientTransport } from '../../core/transport/types.js';
+import { useChatTransport } from './use-chat-transport.js';
+
+/** Options for {@link useMessageSync}. */
+export interface UseMessageSyncOptions {
+  /**
+   * 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.
+   */
+  channelName?: string;
+  /** 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.
- * @param transport - The client transport to observe, or null/undefined if not yet available.
- * @param setMessages - The `setMessages` updater function from useChat.
+ * Subscribe to view updates and sync them into `useChat()`'s overlay.
+ * @param options - Hook options.
+ * @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 = (
-  transport: ClientTransport<unknown, AI.UIMessage> | null | undefined,
-  setMessages: (updater: (prev: AI.UIMessage[]) => AI.UIMessage[]) => void,
-): void => {
+export const useMessageSync = ({ setMessages, channelName, skip }: UseMessageSyncOptions): void => {
+  const { session, chatTransport, chatTransportError } = useChatTransport({ channelName, skip });
+
+  const resolved = !skip && !chatTransportError;
+  const view = resolved ? session.view : undefined;
+  const resolvedChatTransport = resolved ? chatTransport : undefined;
+
+  const [gated, setGated] = useState(false);
+
+  // Subscribe to the ChatTransport's streaming state. Reset on transport
+  // change so a stale `true` doesn't permanently suppress syncs.
   useEffect(() => {
-    if (!transport) return;
-    const unsubscribe = transport.on('message', () => {
-      setMessages(() => transport.getMessages());
-    });
-    return unsubscribe;
-  }, [transport, setMessages]);
+    if (!resolvedChatTransport) {
+      setGated(false);
+      return;
+    }
+    setGated(resolvedChatTransport.streaming);
+    return resolvedChatTransport.onStreamingChange(setGated);
+  }, [resolvedChatTransport]);
+
+  // Subscribe to view updates and sync, unless gated.
+  useEffect(() => {
+    if (!view || gated) return;
+
+    const sync = (): void => {
+      setMessages((overlay) =>
+        mergeMessages(
+          view.getMessages().map((m) => m.message),
+          overlay,
+        ),
+      );
+    };
+
+    // Sync immediately to cover gate-open and initial mount.
+    sync();
+
+    return view.on('update', sync);
+  }, [view, setMessages, gated]);
 };

package/src/vercel/react/vite.config.ts

@@ -19,12 +19,14 @@ export default defineConfig({
       formats: ['es', 'umd'],
     },
     rollupOptions: {
-      external: ['ably', 'ai', 'react'],
+      external: ['ably', 'ably/react', 'react', 'react/jsx-runtime', 'react/jsx-dev-runtime'],
       output: {
         globals: {
           ably: 'Ably',
-          ai: 'AI',
+          'ably/react': 'AblyReact',
           react: 'React',
+          'react/jsx-runtime': 'ReactJsxRuntime',
+          'react/jsx-dev-runtime': 'ReactJsxDevRuntime',
         },
       },
     },

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';
+};