[AIT-681] Compute fork metadata for edits via useChat

demo/vercel/react/use-chat/src/app/chat.tsx

@@ -87,6 +87,7 @@ export function Chat({ chatId, clientId, historyLimit }: { chatId: string; clien
           loading={loading}
           onLoadOlder={loadOlder}
           onRegenerate={(messageId) => regenerate({ messageId })}
+          onEdit={(messageId, newText) => sendMessage({ text: newText, messageId })}
         />
         <InputBar
           onSend={(text) => sendMessage({ text })}

demo/vercel/react/use-chat/src/app/components/message-bubble.tsx

@@ -1,12 +1,14 @@
 'use client';
 
+import { useState } from 'react';
 import type { UIMessage, DynamicToolUIPart } from 'ai';
 import { ToolInvocation } from './tool-invocation';
 
 interface MessageBubbleProps {
   message: UIMessage;
   headers: Record<string, string> | undefined;
   onRegenerate?: () => void;
+  onEdit?: (newText: string) => void;
 }
 
 function Badge({ label, value, color }: { label: string; value: string; color: string }) {
@@ -55,8 +57,22 @@ function bubbleClasses(isUser: boolean, status: string | undefined): string {
   return `${base} bg-zinc-900 text-zinc-300 border border-zinc-800`;
 }
 
-export function MessageBubble({ message, headers, onRegenerate }: MessageBubbleProps) {
+export function MessageBubble({ message, headers, onRegenerate, onEdit }: MessageBubbleProps) {
   const isUser = message.role === 'user';
+  const [editing, setEditing] = useState(false);
+  const textContent = message.parts
+    .filter((p): p is Extract<typeof p, { type: 'text' }> => p.type === 'text')
+    .map((p) => p.text)
+    .join('');
+  const [editText, setEditText] = useState(textContent);
+
+  const handleEditSubmit = (e: React.FormEvent) => {
+    e.preventDefault();
+    const trimmed = editText.trim();
+    if (!trimmed || !onEdit) return;
+    setEditing(false);
+    onEdit(trimmed);
+  };
 
   const role = headers?.['x-ably-role'] ?? message.role;
   const clientId = headers?.['x-ably-turn-client-id'];
@@ -66,23 +82,77 @@ export function MessageBubble({ message, headers, onRegenerate }: MessageBubbleP
   return (
     <div className={`flex ${isUser ? 'justify-end' : 'justify-start'}`}>
       <div className="max-w-[75%]">
-        <div className={bubbleClasses(isUser, status)}>
-          {message.parts.map((part, i) => {
-            if (part.type === 'text') return <span key={i}>{part.text}</span>;
-            if (part.type === 'dynamic-tool')
-              return (
-                <ToolInvocation
-                  key={i}
-                  part={part as DynamicToolUIPart}
-                />
-              );
-            return null;
-          })}
-          {!isUser && status === 'streaming' && (
-            <span className="inline-block w-1.5 h-3.5 ml-0.5 bg-amber-500/60 animate-pulse rounded-sm align-text-bottom" />
-          )}
-        </div>
+        {editing ? (
+          <form
+            onSubmit={handleEditSubmit}
+            className="flex flex-col gap-2"
+          >
+            <textarea
+              value={editText}
+              onChange={(e) => setEditText(e.target.value)}
+              className="rounded-md bg-zinc-900 border border-zinc-600 px-3 py-2 text-sm text-zinc-200 outline-none focus:border-zinc-400 min-w-[300px] resize-y"
+              rows={3}
+              autoFocus
+              onKeyDown={(e) => {
+                if (e.key === 'Escape') {
+                  setEditing(false);
+                  setEditText(textContent);
+                }
+              }}
+            />
+            <div className="flex gap-2 justify-end">
+              <button
+                type="button"
+                onClick={() => {
+                  setEditing(false);
+                  setEditText(textContent);
+                }}
+                className="text-xs text-zinc-500 hover:text-zinc-300 transition-colors rounded bg-zinc-800/60 px-2 py-1"
+              >
+                Cancel
+              </button>
+              <button
+                type="submit"
+                disabled={!editText.trim()}
+                className="text-xs text-zinc-200 hover:bg-zinc-600 transition-colors rounded bg-zinc-700 px-2 py-1 disabled:opacity-40"
+              >
+                Save &amp; send
+              </button>
+            </div>
+          </form>
+        ) : (
+          <div className={bubbleClasses(isUser, status)}>
+            {message.parts.map((part, i) => {
+              if (part.type === 'text') return <span key={i}>{part.text}</span>;
+              if (part.type === 'dynamic-tool')
+                return (
+                  <ToolInvocation
+                    key={i}
+                    part={part as DynamicToolUIPart}
+                  />
+                );
+              return null;
+            })}
+            {!isUser && status === 'streaming' && (
+              <span className="inline-block w-1.5 h-3.5 ml-0.5 bg-amber-500/60 animate-pulse rounded-sm align-text-bottom" />
+            )}
+          </div>
+        )}
         <div className="mt-1 flex items-center gap-1.5 flex-wrap">
+          {/* Edit button (user messages) */}
+          {onEdit && !editing && (
+            <button
+              onClick={() => {
+                setEditText(textContent);
+                setEditing(true);
+              }}
+              className="text-[10px] text-zinc-500 hover:text-zinc-200 transition-colors rounded bg-zinc-800/60 px-1.5 py-0.5"
+              title="Edit message"
+            >
+              edit
+            </button>
+          )}
+
           {/* Regenerate button (assistant messages) */}
           {onRegenerate && status !== 'streaming' && (
             <button

demo/vercel/react/use-chat/src/app/components/message-list.tsx

@@ -11,9 +11,10 @@ interface MessageListProps {
   loading: boolean;
   onLoadOlder: () => void;
   onRegenerate: (messageId: string) => void;
+  onEdit: (messageId: string, newText: string) => void;
 }
 
-export function MessageList({ nodes, hasOlder, loading, onLoadOlder, onRegenerate }: MessageListProps) {
+export function MessageList({ nodes, hasOlder, loading, onLoadOlder, onRegenerate, onEdit }: MessageListProps) {
   const endRef = useRef<HTMLDivElement>(null);
   const scrollRef = useRef<HTMLDivElement>(null);
   const prevLastIdRef = useRef<string | undefined>(undefined);
@@ -61,6 +62,7 @@ export function MessageList({ nodes, hasOlder, loading, onLoadOlder, onRegenerat
           message={message}
           headers={headers}
           onRegenerate={message.role === 'assistant' ? () => onRegenerate(message.id) : undefined}
+          onEdit={message.role === 'user' ? (newText: string) => onEdit(message.id, newText) : undefined}
         />
       ))}
       <div ref={endRef} />

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();