@ably/ai-transport 0.1.0 → 0.2.0

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

@@ -1,10 +1,10 @@
 /**
- * Vercel chat transport: wraps a core ClientTransport to satisfy the
+ * Vercel chat transport: wraps a core ClientSession to satisfy the
  * ChatTransport interface that useChat expects.
  *
- * This is a thin adapter — the real logic lives in the core transport.
+ * This is a thin adapter — the real logic lives in the core client session.
  * The chat transport maps Vercel's sendMessages/reconnectToStream contract
- * to the core transport's send/cancel methods.
+ * to the core session's send/cancel methods.
  *
  * useChat manages message state before calling sendMessages:
  * - submit-message (new): appends the new user message, passes the full array
@@ -12,20 +12,33 @@
  *   passes the truncated array with messageId set
  * - regenerate-message: truncates after the target, passes the truncated array
  *
- * The adapter uses `trigger` to determine the history/messages split:
- * - submit-message: last message is new (publish to channel), rest is history
+ * The adapter uses `(trigger, last-message role)` to determine the
+ * history/messages split:
+ * - submit-message + last message is a user message: that last message is new
+ *   (publish to channel), rest is history. A new submit and an edit both take
+ *   this path — an edit just carries a messageId.
+ * - submit-message + last message is an assistant already in the tree
+ *   (continuation): no new messages, entire array is history
  * - regenerate-message: no new messages, entire array is history
  *
- * When messageId is set (edit or regeneration), the adapter computes fork
- * metadata (forkOf/parent) from the conversation tree so the server can
- * place the response on the correct branch.
+ * For an edit (submit-message with messageId) and for forking off an
+ * unresolved tool call, the adapter computes fork metadata (forkOf/parent)
+ * from the conversation tree so the server can place the response on the
+ * correct branch. Regeneration fork metadata is NOT computed here —
+ * `View.regenerate` derives forkOf/parent from the tree itself.
  */
 
 import * as Ably from 'ably';
 import type * as AI from 'ai';
 
-import type { ClientTransport, CloseOptions, SendOptions } from '../../core/transport/types.js';
+import type { CodecMessage } from '../../core/codec/types.js';
+import type { ActiveRun, ClientSession, SendOptions } from '../../core/transport/types.js';
 import { ErrorCode } from '../../errors.js';
+import { EventEmitter } from '../../event-emitter.js';
+import { LogLevel, makeLogger } from '../../logger.js';
+import type { VercelInput, VercelOutput, VercelProjection } from '../codec/index.js';
+import { UIMessageCodec } from '../codec/index.js';
+import { createRunOutputStream } from './run-output-stream.js';
 
 // ---------------------------------------------------------------------------
 // ChatTransport options
@@ -49,24 +62,36 @@ export interface SendMessagesRequestContext {
   messageId?: string;
   /** Previous messages in the conversation (context for the LLM). */
   history: AI.UIMessage[];
-  /** The new message(s) being sent (to publish to the channel). Empty for regeneration. */
+  /** The new message(s) being sent (to publish to the channel). Empty for regeneration and for continuations (an auto-submit where the last message is an already-tracked assistant). */
   messages: AI.UIMessage[];
-  /** The msg-id of the message being forked (regenerated or edited). */
+  /** The codec-message-id of the message being forked — the edited user message, or the preceding assistant when forking off an unresolved tool call. Undefined for regeneration (View.regenerate derives it) and fresh sends. */
   forkOf?: string;
-  /** The msg-id of the predecessor in the conversation thread. */
+  /** The codec-message-id of the predecessor in the conversation thread. */
   parent?: string;
 }
 
+/** Default agent endpoint the transport POSTs invocations to — mirrors Vercel's DefaultChatTransport. */
+export const DEFAULT_VERCEL_API = '/api/chat';
+
 /** Options for customizing the ChatTransport behavior. */
 export interface ChatTransportOptions {
   /**
-   * Customize the POST body before sending. Called by sendMessages()
-   * with the conversation context. Return the body and headers for
-   * the HTTP POST.
-   *
-   * Default: sends all previous messages as `history` in the body.
+   * Endpoint the transport POSTs the invocation pointer to, to wake the
+   * agent. Mirrors useChat's request-driven contract. Default `/api/chat`.
+   */
+  api?: string;
+  /** Fetch credentials mode for the invocation POST (e.g. `'include'`). */
+  credentials?: RequestCredentials;
+  /** Custom fetch implementation for the invocation POST. Defaults to `globalThis.fetch`. */
+  fetch?: typeof globalThis.fetch;
+  /**
+   * Customize the invocation POST before sending. Called by sendMessages()
+   * with the conversation context; the returned `body` is merged into the
+   * POST body (the run's invocation identifiers always take precedence) and
+   * `headers` are added to the request. Use it for auth headers or extra
+   * agent metadata.
    * @param context - The conversation context for the current request.
-   * @returns The body and headers to use for the HTTP POST.
+   * @returns The body and headers to merge into the invocation POST.
    */
   prepareSendMessagesRequest?: (context: SendMessagesRequestContext) => {
     body?: Record<string, unknown>;
@@ -131,9 +156,9 @@ export interface ChatTransport {
   ) => Promise<ReadableStream<AI.UIMessageChunk> | null>;
 
   /** Close the underlying transport, releasing all resources. */
-  close(options?: CloseOptions): Promise<void>;
+  close(): Promise<void>;
 
-  /** Whether an own-turn stream is currently being consumed by useChat. */
+  /** Whether an own-run stream is currently being consumed by useChat. */
   readonly streaming: boolean;
 
   /**
@@ -153,11 +178,17 @@ export interface ChatTransport {
 /**
  * Wrap a ReadableStream in a passthrough TransformStream that resolves a
  * promise when the stream completes or errors. The returned stream passes
- * all chunks through unchanged.
+ * all chunks through unchanged, and `fail(reason)` errors the readable side
+ * useChat consumes without cancelling or otherwise disturbing the source run
+ * stream (used when the agent-invocation POST fails).
  * @param source - The original stream to wrap.
- * @returns The wrapped stream and a `done` promise that resolves when the stream closes.
+ * @returns The wrapped stream, a `done` promise that resolves when the stream
+ *   closes, and a `fail` callback that errors the wrapped stream.
  */
-const wrapStreamWithDone = <T>(source: ReadableStream<T>): { stream: ReadableStream<T>; done: Promise<void> } => {
+
+const wrapStreamWithDone = <T>(
+  source: ReadableStream<T>,
+): { stream: ReadableStream<T>; done: Promise<void>; fail: (reason: Ably.ErrorInfo) => void } => {
   let resolveDone: () => void;
   const done = new Promise<void>((resolve) => {
     resolveDone = resolve;
@@ -169,28 +200,51 @@ const wrapStreamWithDone = <T>(source: ReadableStream<T>): { stream: ReadableStr
     },
   });
 
-  // Pipe in the background. If the source errors or is cancelled, resolve
-  // done so the serialization queue advances.
+  // Aborting this signal errors the destination (the readable useChat reads)
+  // with the abort reason. `preventCancel` keeps the source run stream intact
+  // so the tree/observers are unaffected — only the useChat-facing view fails.
+  const failController = new AbortController();
+
+  // Pipe in the background. If the source errors/cancels, or `fail()` aborts,
+  // resolve done so the serialization queue advances.
   // Fire-and-forget: the pipe runs independently; errors surface through
   // the readable side that useChat consumes.
-  source.pipeTo(passthrough.writable).catch(() => {
+  source.pipeTo(passthrough.writable, { signal: failController.signal, preventCancel: true }).catch(() => {
     resolveDone();
   });
 
-  return { stream: passthrough.readable, done };
+  return {
+    stream: passthrough.readable,
+    done,
+    fail: (reason: Ably.ErrorInfo) => {
+      failController.abort(reason);
+    },
+  };
 };
 
 // ---------------------------------------------------------------------------
 // Unresolved tool call detection
 // ---------------------------------------------------------------------------
 
+/**
+ * Whether a UIMessage part is a tool part — either the codec-normalised
+ * `dynamic-tool` shape or the AI SDK's statically-declared `tool-${name}`
+ * shape. Both carry `toolCallId` and `state`; the shape check at the end
+ * is defensive against a future AI SDK release introducing a non-tool
+ * variant under the `tool-` prefix (none exists today).
+ * @param part - The UIMessage part to inspect.
+ * @returns True when the part is a tool part of either representation.
+ */
+const _isToolPart = (part: AI.UIMessage['parts'][number]): part is AI.DynamicToolUIPart | AI.ToolUIPart =>
+  (part.type === 'dynamic-tool' || part.type.startsWith('tool-')) && 'toolCallId' in part && 'state' in part;
+
 /**
  * Whether an assistant message has a `dynamic-tool` part that can't resolve
  * without further user action. Matches:
  * - `input-streaming` / `input-available` — tool call emitted, not yet run.
  * - `approval-requested` — waiting for the user.
  *
- * Excludes `approval-responded` (streamText will run the tool this turn)
+ * Excludes `approval-responded` (streamText will run the tool this run)
  * and all terminal `output-*` states.
  * @param msg - The UIMessage to inspect.
  * @returns True when a fork-on-send is warranted to avoid shipping a
@@ -200,55 +254,198 @@ const hasUnresolvedToolCall = (msg: AI.UIMessage): boolean =>
   msg.role === 'assistant' &&
   msg.parts.some(
     (p) =>
-      p.type === 'dynamic-tool' &&
+      _isToolPart(p) &&
       (p.state === 'input-streaming' || p.state === 'input-available' || p.state === 'approval-requested'),
   );
 
+/**
+ * `dynamic-tool` part states that mean "the LLM produced a tool call and
+ * is waiting on it". Used to detect new client-side resolutions in the
+ * useChat overlay relative to the tree.
+ */
+const UNRESOLVED_TOOL_STATES = new Set(['input-streaming', 'input-available', 'approval-requested']);
+
+/**
+ * Walk the useChat message overlay against the session tree and synthesize
+ * the {@link VercelInput}s needed to resolve every `dynamic-tool` part the
+ * user acted on (executed a tool, approved, denied) but the tree's reduced
+ * state hasn't reflected yet.
+ *
+ * Each input carries the prior assistant's tree codec-message-id (the one
+ * holding the original `dynamic-tool` part the resolution targets) in its
+ * `codecMessageId` field, so the encoder stamps `codec-message-id`
+ * and the reducer's direct-fold path lands the resolution on that assistant
+ * in one step — no cross-message redirect-by-toolCallId fallback. Every
+ * variant rides the `ai-input` wire, matching its publisher (client → input).
+ *
+ * The resulting inputs are passed alongside the continuation `view.send`
+ * so the channel publish and the continuation POST land as ONE atomic
+ * operation — the agent's `loadProjection()` history fetch is guaranteed
+ * to see them because the channel publish happens before the POST inside
+ * `_internalSend`.
+ *
+ * Three resolutions are produced:
+ *
+ * - `approval-responded` overlay vs `approval-requested` tree →
+ *   `tool-approval-response` carrying the user's decision
+ *   (`approved` = `overlayPart.approval.approved`, i.e. approve or deny)
+ * - `output-available` overlay vs unresolved tree → `tool-result`
+ * - `output-error` overlay vs unresolved tree → `tool-result-error`
+ * @param codecMessages - The visible tree messages paired with their codec-message-ids.
+ * @param messages - useChat's local overlay messages.
+ * @returns The continuation inputs to publish, in tree order. Each input
+ *   carries its own `codecMessageId` targeting the prior assistant it folds
+ *   onto.
+ */
+const deriveContinuationInputs = (
+  codecMessages: CodecMessage<AI.UIMessage>[],
+  messages: AI.UIMessage[],
+): VercelInput[] => {
+  const inputs: VercelInput[] = [];
+  for (const overlay of messages) {
+    if (overlay.role !== 'assistant') continue;
+    // Match the overlay to its tree message by domain id (both sides
+    // reconstruct the same stream id), but address the emitted inputs by
+    // the tree message's codec-message-id — the agent folds tool
+    // resolutions onto the assistant by codec-message-id, never by the
+    // domain `message.id`.
+    const treeEntry = codecMessages.find((p) => p.message.id === overlay.id);
+    if (!treeEntry) continue;
+    const { codecMessageId, message: treeMessage } = treeEntry;
+
+    for (const overlayPart of overlay.parts) {
+      if (!_isToolPart(overlayPart)) continue;
+      // The codec normalises every tool part to `dynamic-tool`, but the
+      // AI SDK's useChat overlay emits `tool-${name}` parts for statically
+      // declared tools. Match by toolCallId rather than the type prefix
+      // so the cross-representation comparison works regardless of which
+      // side the tool was declared on.
+      const treePart = treeMessage.parts.find(
+        (p: AI.UIMessage['parts'][number]): p is AI.DynamicToolUIPart | AI.ToolUIPart =>
+          _isToolPart(p) && p.toolCallId === overlayPart.toolCallId,
+      );
+
+      // Approval response: useChat's `addToolApprovalResponse` flipped the
+      // overlay part to `approval-responded` while the tree still sits on
+      // `approval-requested`. Publish a `tool-approval-response` TInput so the
+      // agent's projection sees the decision.
+      if (overlayPart.state === 'approval-responded' && (!treePart || treePart.state === 'approval-requested')) {
+        inputs.push(
+          UIMessageCodec.createToolApprovalResponse(codecMessageId, {
+            toolCallId: overlayPart.toolCallId,
+            approved: overlayPart.approval.approved,
+            ...(overlayPart.approval.reason === undefined ? {} : { reason: overlayPart.approval.reason }),
+          }),
+        );
+        continue;
+      }
+
+      // Client-tool resolution: overlay has `output-available` / `output-error`
+      // while the tree's part is still unresolved. Construct a TInput
+      // variant (not a UIMessageChunk) so the encoder publishes on the
+      // `ai-input` wire — client tool results belong on `ai-input`, matching
+      // their client publisher, not on `ai-output`.
+      if (overlayPart.state !== 'output-available' && overlayPart.state !== 'output-error') continue;
+      // Tree already resolved (echo arrived back) — nothing to do.
+      if (treePart && !UNRESOLVED_TOOL_STATES.has(treePart.state)) continue;
+
+      if (overlayPart.state === 'output-available') {
+        inputs.push(
+          UIMessageCodec.createToolResult(codecMessageId, {
+            toolCallId: overlayPart.toolCallId,
+            output: overlayPart.output,
+          }),
+        );
+      } else {
+        inputs.push(
+          UIMessageCodec.createToolResultError(codecMessageId, {
+            toolCallId: overlayPart.toolCallId,
+            message: overlayPart.errorText,
+          }),
+        );
+      }
+    }
+  }
+  return inputs;
+};
+
+/**
+ * Find the codec-message-id immediately preceding the message identified by
+ * domain id `domainId` in the flat visible conversation. The target is
+ * located by its domain `message.id` (the id useChat references), but the
+ * returned value is the predecessor's codec-message-id — never a domain id.
+ * Returns undefined if the target is the first message or not found.
+ * @param codecMessages - Visible messages paired with their codec-message-ids.
+ * @param domainId - The domain id of the target message.
+ * @returns The predecessor's codec-message-id, or undefined.
+ */
+const findPredecessorCodecId = (codecMessages: CodecMessage<AI.UIMessage>[], domainId: string): string | undefined => {
+  const idx = codecMessages.findIndex((p) => p.message.id === domainId);
+  if (idx <= 0) return undefined;
+  return codecMessages[idx - 1]?.codecMessageId;
+};
+
 // ---------------------------------------------------------------------------
 // Factory
 // ---------------------------------------------------------------------------
 
+/** Internal EventEmitter events map backing the transport's streaming state. */
+interface ChatTransportEventsMap {
+  /** Fired on every streaming-state transition with the new value. */
+  streaming: boolean;
+}
+
 /**
- * Create a Vercel ChatTransport from a core ClientTransport.
+ * Create a Vercel ChatTransport from a core ClientSession.
  *
  * Exposes a `streaming` flag and `onStreamingChange` callback so that
- * `useMessageSync` can gate `setMessages` calls during active own-turn
+ * `useMessageSync` can gate `setMessages` calls during active own-run
  * streams, preventing the push/replace ID mismatch in useChat's `write()`.
  *
  * Note: concurrent `sendMessage` calls from the same user are a useChat
  * limitation that cannot be fixed from the transport layer. The
  * developer must respect useChat's `status` and only call `sendMessage`
  * when status is `'ready'`.
- * @param transport - The core client transport to wrap.
+ * @param session - The core client session to wrap.
  * @param chatOptions - Optional hooks for customizing request construction.
  * @returns A {@link ChatTransport} compatible with Vercel's useChat hook.
  */
 export const createChatTransport = (
-  transport: ClientTransport<AI.UIMessageChunk, AI.UIMessage>,
+  session: ClientSession<VercelInput, VercelOutput, VercelProjection, AI.UIMessage>,
   chatOptions?: ChatTransportOptions,
 ): ChatTransport => {
+  // -- Invocation POST config (the transport owns waking the agent) ----------
+  const api = chatOptions?.api ?? DEFAULT_VERCEL_API;
+  const fetchFn = chatOptions?.fetch ?? globalThis.fetch.bind(globalThis);
+  const credentials = chatOptions?.credentials;
+
   // -- Streaming state -------------------------------------------------------
+  // Backed by the shared EventEmitter for listener error isolation (one bad
+  // onStreamingChange handler can't prevent others from firing or block the
+  // state transition) and uniform emitter behaviour across the SDK. The
+  // factory takes no logger, so a silent one is used — listener exceptions are
+  // swallowed by the emitter rather than surfaced.
   let _streaming = false;
-  const streamingCallbacks = new Set<(streaming: boolean) => void>();
+  const emitter = new EventEmitter<ChatTransportEventsMap>(makeLogger({ logLevel: LogLevel.Silent }));
 
   const setStreaming = (value: boolean): void => {
     _streaming = value;
-    for (const cb of streamingCallbacks) {
-      try {
-        cb(value);
-      } catch {
-        // Isolate subscriber errors so one bad handler doesn't prevent
-        // other subscribers from being notified or block the streaming
-        // state transition.
-      }
-    }
+    emitter.emit('streaming', value);
   };
 
   // -- sendMessages implementation -------------------------------------------
 
   const sendMessages: ChatTransport['sendMessages'] = async (opts) => {
     const { messages, abortSignal, trigger, messageId } = opts;
-    const allNodes = transport.view.flattenNodes();
+
+    // The visible messages paired with their codec-message-ids. useChat
+    // references messages by their domain `message.id`; we match on that to
+    // locate a message in the tree, then route every transport operation by
+    // the message's codec-message-id (the SDK never correlates on the domain
+    // id, which may differ from the codec-message-id).
+    const codecMessages = session.view.getMessages();
+    const codecIdByDomainId = new Map(codecMessages.map((m) => [m.message.id, m.codecMessageId]));
+    const codecIdOf = (domainId: string): string | undefined => codecIdByDomainId.get(domainId);
 
     // useChat calls sendMessages in three distinct modes. We disambiguate
     // by (trigger, last-message role) so each mode dispatches correctly:
@@ -265,27 +462,29 @@ export const createChatTransport = (
     // Continuation mode must NOT publish the assistant as a new message or
     // treat messageId as a fork target — useChat v6's sendAutomaticallyWhen
     // path always sets messageId to the last message id regardless.
-    //
-    // Client-side tool outputs are expected to be staged on the transport
-    // via transport.stageEvents() before this runs; the core transport
-    // flushes staged events into the POST body automatically.
     const lastMessage = messages.at(-1);
-    const lastMessageNode = lastMessage ? allNodes.find((n) => n.message.id === lastMessage.id) : undefined;
-    const isContinuation = trigger === 'submit-message' && lastMessage?.role === 'assistant' && !!lastMessageNode;
+    const lastMessageInTree = !!lastMessage && codecIdByDomainId.has(lastMessage.id);
+    const isContinuation = trigger === 'submit-message' && lastMessage?.role === 'assistant' && lastMessageInTree;
 
     // Fork-on-unresolved-tool: user sent a new message while the preceding
     // assistant has an unresolved tool call (approval-requested, input-*).
     // Fork the new message off the preceding assistant so the unresolved
-    // tool call stays dormant on a sibling branch. Inference this turn runs
+    // tool call stays dormant on a sibling branch. Inference for this run runs
     // on the clean fork — the LLM never sees the dangling tool_use.
     //
     // Only applies to fresh user-message submits (not continuations, not
     // regenerates, not edits-with-messageId).
+    //
+    // `messages.at(-1)` is the fresh user-prompt being submitted right now;
+    // `messages.at(-2)` is therefore the prior assistant whose tool state
+    // we need to inspect for the unresolved-tool gate below.
     const precedingMessage =
       trigger === 'submit-message' && !messageId && lastMessage?.role === 'user' ? messages.at(-2) : undefined;
-    const forkSource =
-      precedingMessage && hasUnresolvedToolCall(precedingMessage)
-        ? allNodes.find((n) => n.message.id === precedingMessage.id)
+    // The domain id of the preceding assistant when it carries an unresolved
+    // tool call and is present in the tree — the new user message forks off it.
+    const forkSourceDomainId =
+      precedingMessage && hasUnresolvedToolCall(precedingMessage) && codecIdByDomainId.has(precedingMessage.id)
+        ? precedingMessage.id
         : undefined;
 
     // Determine the history/messages split based on mode.
@@ -309,35 +508,27 @@ export const createChatTransport = (
       // When forking off an unresolved tool call, drop the unresolved
       // assistant from history too — it belongs on the sibling branch, not
       // the ancestor chain of the new message.
-      history = forkSource ? messages.slice(0, -2) : messages.slice(0, -1);
+      history = forkSourceDomainId ? messages.slice(0, -2) : messages.slice(0, -1);
     }
 
-    // Compute fork metadata. Only set in regenerate or edit modes — in
-    // continuation mode we do NOT fork, we continue the branch.
+    // Compute fork metadata for edit (submit-message with messageId) and
+    // fork-on-unresolved-tool. Regenerate is NOT precomputed here —
+    // `View.regenerate` derives forkOf/parent from the tree itself and
+    // overrides anything we'd set.
     let forkOf: string | undefined;
     let parent: string | undefined;
 
-    if (messageId && !isContinuation) {
-      // Regeneration: messageId = assistant to regenerate.
-      // Edit (submit-message with user message and messageId): messageId = user being replaced.
-      // In both cases forkOf = the x-ably-msg-id, parent = that message's parent.
-      forkOf = messageId;
-      const node = allNodes.find((n) => n.message.id === messageId);
-      if (node) {
-        forkOf = node.msgId;
-        parent = node.parentId;
-      }
-    } else if (isContinuation) {
-      // Continuation: the server's next assistant message is a child of the
-      // last assistant (no fork). Pass parent so the server places the new
-      // message correctly in the tree. isContinuation narrows lastMessageNode
-      // to defined.
-      parent = lastMessageNode.msgId;
-    } else if (forkSource) {
+    if (trigger === 'submit-message' && messageId && !isContinuation) {
+      // Edit: messageId is the domain id of the user message being replaced.
+      // forkOf = its codec-message-id, parent = the immediately-preceding
+      // codec-message-id in the flat conversation.
+      forkOf = codecIdOf(messageId);
+      parent = findPredecessorCodecId(codecMessages, messageId);
+    } else if (forkSourceDomainId) {
       // Fork off the preceding assistant — the new user message becomes a
       // sibling of the unresolved tool call assistant, rooted at its parent.
-      forkOf = forkSource.msgId;
-      parent = forkSource.parentId;
+      forkOf = codecIdOf(forkSourceDomainId);
+      parent = findPredecessorCodecId(codecMessages, forkSourceDomainId);
     }
 
     let sendBody: Record<string, unknown>;
@@ -356,38 +547,102 @@ export const createChatTransport = (
       sendBody = prepared.body ?? {};
       sendHeaders = prepared.headers;
     } else {
-      const historyIds = new Set(history.map((m) => m.id));
-      const historyNodes = allNodes.filter((n) => historyIds.has(n.message.id));
-      sendBody = {
-        history: historyNodes,
-        chatId: opts.chatId,
-        trigger,
-        ...(messageId !== undefined && { messageId }),
-        ...(forkOf !== undefined && { forkOf }),
-        ...(parent !== undefined && { parent }),
-      };
+      sendBody = {};
       sendHeaders = undefined;
     }
 
-    const sendOpts: SendOptions = { body: sendBody, headers: sendHeaders };
+    const sendOpts: SendOptions = {};
     if (forkOf !== undefined) sendOpts.forkOf = forkOf;
     if (parent !== undefined) sendOpts.parent = parent;
+    // Continuations reuse the suspended assistant's runId so the agent's
+    // existing run resumes under a fresh invocation rather than spinning
+    // up a brand-new run. `isContinuation` implies `lastMessage` is defined.
+    if (isContinuation) {
+      // `isContinuation` implies `lastMessage` is defined (it gates on
+      // `lastMessage?.role`). Route the runId lookup by codec-message-id.
+      const codecId = codecIdOf(lastMessage.id);
+      const run = codecId === undefined ? undefined : session.view.runOf(codecId);
+      if (run) sendOpts.runId = run.runId;
+    }
 
-    // A single dispatch path: view.send with the (possibly empty)
-    // newMessages array. Any events staged via transport.stageEvents()
-    // flow automatically through _internalSend into the POST body.
-    const turn = await transport.view.send(newMessages, sendOpts);
+    // Dispatch by mode:
+    //
+    // - Continuation: derive tool-resolution events from useChat's overlay
+    //   vs the tree and pair each with the prior assistant's tree codec-message-id —
+    //   the SDK stamps the wire's `codec-message-id` to that id so the
+    //   reducer's direct fold path runs (no redirect, no consume).
+    // - Regenerate: route through `view.regenerate`. The View mints a
+    //   wire-only regenerate event (`ait-regenerate`) carrying
+    //   `forkOf=A1` / `parent=U1` on transport headers. U1 is NOT
+    //   republished — A1 and A2 group as tree siblings under U1 via the
+    //   existing forkOf machinery. The LLM receives the truncated history
+    //   through U1 inclusive via the body.
+    // - Fresh send / edit: publish the new user-message input(s) via
+    //   `view.send`.
+    let run: ActiveRun;
+    if (isContinuation) {
+      const inputs = deriveContinuationInputs(codecMessages, messages);
+      run = await session.view.send(inputs, sendOpts);
+    } else if (trigger === 'regenerate-message') {
+      if (messageId === undefined) {
+        throw new Ably.ErrorInfo(
+          'unable to regenerate; regenerate-message trigger fired without messageId',
+          ErrorCode.InvalidArgument,
+          400,
+        );
+      }
+      // useChat passes the assistant's domain id; route by its codec-message-id.
+      const regenCodecId = codecIdOf(messageId);
+      if (regenCodecId === undefined) {
+        throw new Ably.ErrorInfo(
+          `unable to regenerate; message not visible: ${messageId}`,
+          ErrorCode.InvalidArgument,
+          400,
+        );
+      }
+      run = await session.view.regenerate(regenCodecId, sendOpts);
+    } else {
+      const inputs = newMessages.map((m) => UIMessageCodec.createUserMessage(m));
+      run = await session.view.send(inputs, sendOpts);
+    }
+
+    // Build the consumer-facing stream from the Tree's events for this run.
+    // Streaming is a useChat concern owned by the Vercel layer; the core
+    // session exposes no per-run stream. Key it on
+    // `run.inputCodecMessageId` — the triggering input's codec-message-id, which
+    // the client owns from send time and the agent echoes as
+    // `input-codec-message-id`. The agent mints the runId, supplied as
+    // `run.runId` (a promise) for the run-end safety-net.
+    const runStream = createRunOutputStream(session, run.runId, run.inputCodecMessageId);
 
     if (abortSignal) {
-      abortSignal.addEventListener('abort', () => void transport.cancel({ all: true }), {
-        once: true,
-      });
+      const onAbort = (): void => {
+        // Best-effort cancel via the run handle (knows its own key / runId);
+        // the core resolves the runId once the agent mints it.
+        void run.cancel();
+        // Close the consumer stream immediately so useChat's reader ends
+        // without waiting for the agent's run-end round-trip.
+        runStream.close();
+      };
+      // useChat sets `status: 'submitted'` synchronously inside `makeRequest`
+      // BEFORE awaiting `transport.sendMessages`. That immediately enables
+      // the Stop button in the UI. If the user clicks Stop while
+      // `session.view.send` is still awaiting the run-start ack (which
+      // can take seconds for a real LLM), useChat aborts the signal before
+      // we ever get here. `addEventListener('abort', ...)` does not fire
+      // for an already-aborted signal, so we'd silently lose the cancel
+      // and the agent would keep streaming.
+      if (abortSignal.aborted) {
+        onAbort();
+      } else {
+        abortSignal.addEventListener('abort', onAbort, { once: true });
+      }
     }
 
     // Wrap the stream to detect completion. The streaming flag gates
     // useMessageSync so that setMessages doesn't interfere with
     // useChat's internal write() during active streams.
-    const { stream, done } = wrapStreamWithDone(turn.stream);
+    const { stream, done, fail } = wrapStreamWithDone(runStream.stream);
     setStreaming(true);
 
     // Fire-and-forget: clear the streaming flag when the stream ends.
@@ -395,6 +650,45 @@ export const createChatTransport = (
       setStreaming(false);
     });
 
+    // Wake the agent: POST the invocation pointer to the configured endpoint.
+    // useChat's transport contract is request-driven, so the transport owns
+    // this POST (the core session is HTTP-free). Fire-and-forget — `await`
+    // would delay the stream return, and the agent's response arrives over
+    // the Ably channel, not the HTTP response. The run's invocation
+    // identifiers always win over any custom body so the agent can parse it
+    // via Invocation.fromJSON. A failed POST means the agent never woke, so
+    // error the useChat-facing stream; the core run and observers are
+    // untouched.
+    const postBody = { ...sendBody, ...run.toInvocation().toJSON() };
+    fetchFn(api, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json', ...sendHeaders },
+      body: JSON.stringify(postBody),
+      ...(credentials ? { credentials } : {}),
+    })
+      .then((response) => {
+        if (!response.ok) {
+          fail(
+            new Ably.ErrorInfo(
+              `unable to send; HTTP POST to ${api} returned ${String(response.status)} ${response.statusText}`,
+              ErrorCode.SessionSendFailed,
+              response.status,
+            ),
+          );
+        }
+      })
+      .catch((error: unknown) => {
+        const cause = error instanceof Ably.ErrorInfo ? error : undefined;
+        fail(
+          new Ably.ErrorInfo(
+            `unable to send; HTTP POST to ${api} failed: ${error instanceof Error ? error.message : String(error)}`,
+            ErrorCode.SessionSendFailed,
+            500,
+            cause,
+          ),
+        );
+      });
+
     return stream;
   };
 
@@ -408,16 +702,16 @@ export const createChatTransport = (
     // eslint-disable-next-line unicorn/no-null, @typescript-eslint/promise-function-async -- null is required by the AI SDK ChatTransport contract; no await needed
     reconnectToStream: () => Promise.resolve(null),
 
-    close: async (options?: CloseOptions) => transport.close(options),
+    close: async () => session.close(),
 
     get streaming(): boolean {
       return _streaming;
     },
 
     onStreamingChange: (callback: (streaming: boolean) => void): (() => void) => {
-      streamingCallbacks.add(callback);
+      emitter.on('streaming', callback);
       return () => {
-        streamingCallbacks.delete(callback);
+        emitter.off('streaming', callback);
       };
     },
   };