transport/vercel: compute fork metadata for edits via useChat

Our SDK has two forking operations: edit (replace a user message)
and regenerate (re-run an assistant message). The AI SDK's useChat
exposes the same two: sendMessage with a messageId (edit) [1] and
regenerate(). We already made our regenerate compatible with
useChat's — ChatTransport computes forkOf/parent from the
conversation tree when trigger is 'regenerate-message'.

It turns out we missed the edit case. Without fork metadata,
edits via useChat were treated as new messages rather than forks
in the conversation tree.

We fix it by always including fork metadata when the messageId is
provided, regardless of the trigger, thus continuing to cover the
regeneration case and now also the edit case.

[1] https://ai-sdk.dev/docs/reference/ai-sdk-ui/use-chat#send-message

[AIT-681]

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: lawrence-forooghian

src/vercel/transport/chat-transport.ts

@@ -7,12 +7,18 @@
  * to the core transport's send/cancel methods.
  *
  * useChat manages message state before calling sendMessages:
- * - submit-message: appends the new user message, passes the full array
+ * - submit-message (new): appends the new user message, passes the full array
+ * - submit-message (edit): truncates after the edited message, replaces it,
+ *   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
  * - 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.
  */
 
 import * as Ably from 'ably';
@@ -35,8 +41,10 @@ export interface SendMessagesRequestContext {
   /** What triggered the request: user sent a message, or requested regeneration. */
   trigger: 'submit-message' | 'regenerate-message';
   /**
-   * The message ID for regeneration requests. Identifies which assistant
-   * message to regenerate. Undefined for submit-message.
+   * The message ID for edit or regeneration requests. For regeneration,
+   * identifies the assistant message to regenerate. For edits (submit-message
+   * with messageId), identifies the user message being replaced. Undefined
+   * when submitting a new message.
    */
   messageId?: string;
   /** Previous messages in the conversation (context for the LLM). */
@@ -239,11 +247,19 @@ export const createChatTransport = (
     }
 
     // Compute fork metadata from the conversation tree.
+    // For regeneration: messageId is the assistant message being regenerated.
+    // For edit: messageId is the user message being replaced.
+    // In both cases: forkOf = the x-ably-msg-id of that message,
+    //   parent = the parent of that message in the tree.
     let forkOf: string | undefined;
     let parent: string | undefined;
 
-    if (trigger === 'regenerate-message' && messageId) {
+    if (messageId) {
       forkOf = messageId;
+      // Look up the message in the tree to resolve x-ably-msg-id.
+      // messageId comes from useChat (UIMessage.id) — scan the flattened
+      // nodes to find the one whose domain message matches this ID.
+      // Uses the transport's default view — ChatTransport is single-view (one useChat per channel).
       const node = transport.view.flattenNodes().find((n) => n.message.id === messageId);
       if (node) {
         forkOf = node.msgId;

test/vercel/transport/chat-transport.test.ts

@@ -278,6 +278,97 @@ describe('createChatTransport', () => {
     });
   });
 
+  describe('sendMessages — submit-message with messageId (edit)', () => {
+    it('resolves fork metadata from the conversation tree', async () => {
+      const { transport, send, view, mockTurn } = createMockTransport();
+
+      const edited = makeMessage('ui-msg-id');
+      (view.flattenNodes as ReturnType<typeof vi.fn>).mockReturnValue([
+        {
+          message: edited,
+          msgId: 'wire-msg-id',
+          parentId: 'wire-parent-id',
+          forkOf: undefined,
+          headers: {},
+          serial: undefined,
+        },
+      ]);
+
+      const chat = createChatTransport(transport);
+
+      const streamPromise = chat.sendMessages({
+        trigger: 'submit-message',
+        chatId: 'chat-1',
+        messageId: 'ui-msg-id',
+        messages: [edited],
+        abortSignal: undefined,
+      });
+
+      mockTurn.close();
+      await streamPromise;
+
+      const [, opts] = send.mock.calls[0] as [AI.UIMessage[], SendOptions];
+      expect(opts.forkOf).toBe('wire-msg-id');
+      expect(opts.parent).toBe('wire-parent-id');
+    });
+
+    it('falls back to raw messageId when node not found in tree', async () => {
+      const { transport, send, view, mockTurn } = createMockTransport();
+      (view.flattenNodes as ReturnType<typeof vi.fn>).mockReturnValue([]);
+
+      const chat = createChatTransport(transport);
+
+      const streamPromise = chat.sendMessages({
+        trigger: 'submit-message',
+        chatId: 'chat-1',
+        messageId: 'unknown-id',
+        messages: [makeMessage('1')],
+        abortSignal: undefined,
+      });
+
+      mockTurn.close();
+      await streamPromise;
+
+      const [, opts] = send.mock.calls[0] as [AI.UIMessage[], SendOptions];
+      expect(opts.forkOf).toBe('unknown-id');
+      expect(opts.parent).toBeUndefined();
+    });
+
+    it('sends the edited message as new and prior messages as history', async () => {
+      const { transport, send, view, mockTurn } = createMockTransport();
+
+      const m1 = makeMessage('1');
+      const edited = makeMessage('2');
+
+      (view.flattenNodes as ReturnType<typeof vi.fn>).mockReturnValue([
+        { message: m1, msgId: 'n1', parentId: undefined, forkOf: undefined, headers: {}, serial: undefined },
+        { message: edited, msgId: 'n2', parentId: 'n1', forkOf: undefined, headers: {}, serial: undefined },
+      ]);
+
+      const chat = createChatTransport(transport);
+
+      const streamPromise = chat.sendMessages({
+        trigger: 'submit-message',
+        chatId: 'chat-1',
+        messageId: '2',
+        messages: [m1, edited],
+        abortSignal: undefined,
+      });
+
+      mockTurn.close();
+      await streamPromise;
+
+      const [msgs, opts] = send.mock.calls[0] as [AI.UIMessage[], SendOptions];
+      expect(msgs).toEqual([edited]);
+      // CAST: body is always set by the adapter; narrowing to non-undefined.
+      // eslint-disable-next-line @typescript-eslint/non-nullable-type-assertion-style -- prefer `as` over `!` per TYPES.md
+      const body = opts.body as Record<string, unknown>;
+      const bodyHistory = body.history as { message: AI.UIMessage }[];
+      expect(bodyHistory).toHaveLength(1);
+      expect(bodyHistory.at(0)?.message).toEqual(m1);
+    });
+  });
+
   describe('real stream return', () => {
     it('returns the turn stream with chunks flowing through', async () => {
       const { transport, mockTurn } = createMockTransport();