Skip to content

Commit ddd54ba

Browse files
lawrence-forooghianclaude
lawrence-forooghian
and
claude
committed
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. Also update the use-chat demo to demonstrate editing (it already demonstrates regeneration), using the same UX as the use-client-transport demo. [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>
1 parent f3112f7 commit ddd54ba

File tree

5 files changed

+260
-58
lines changed

5 files changed

+260
-58
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ export function Chat({ chatId, clientId, historyLimit }: { chatId: string; clien
8888
loading={loading}
8989
onLoadOlder={loadOlder}
9090
onRegenerate={(messageId) => regenerate({ messageId })}
91+
onEdit={(messageId, text) => sendMessage({ text, messageId })}
9192
/>
9293
<InputBar
9394
onSend={(text) => sendMessage({ text })}

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

Lines changed: 145 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
'use client';
22

3+
import { useState } from 'react';
34
import type { UIMessage, DynamicToolUIPart } from 'ai';
45
import { ToolInvocation } from './tool-invocation';
56

67
interface MessageBubbleProps {
78
message: UIMessage;
89
headers: Record<string, string> | undefined;
910
onRegenerate?: () => void;
11+
onEdit?: (newText: string) => void;
1012
}
1113

1214
function Badge({ label, value, color }: { label: string; value: string; color: string }) {
@@ -55,71 +57,161 @@ function bubbleClasses(isUser: boolean, status: string | undefined): string {
5557
return `${base} bg-zinc-900 text-zinc-300 border border-zinc-800`;
5658
}
5759

58-
export function MessageBubble({ message, headers, onRegenerate }: MessageBubbleProps) {
60+
// ---------------------------------------------------------------------------
61+
// Inline edit form
62+
// ---------------------------------------------------------------------------
63+
64+
function EditForm({
65+
initialText,
66+
onSubmit,
67+
onCancel,
68+
}: {
69+
initialText: string;
70+
onSubmit: (text: string) => void;
71+
onCancel: () => void;
72+
}) {
73+
const [text, setText] = useState(initialText);
74+
75+
const handleSubmit = (e: React.FormEvent) => {
76+
e.preventDefault();
77+
const trimmed = text.trim();
78+
if (trimmed && trimmed !== initialText) {
79+
onSubmit(trimmed);
80+
}
81+
onCancel();
82+
};
83+
84+
return (
85+
<form
86+
onSubmit={handleSubmit}
87+
className="w-full"
88+
>
89+
<textarea
90+
value={text}
91+
onChange={(e) => setText(e.target.value)}
92+
className="w-full rounded-lg bg-zinc-800 border border-zinc-600 px-3 py-2 text-sm text-zinc-200 outline-none focus:border-zinc-400 resize-none"
93+
rows={Math.min(6, text.split('\n').length + 1)}
94+
autoFocus
95+
onKeyDown={(e) => {
96+
if (e.key === 'Escape') onCancel();
97+
if (e.key === 'Enter' && !e.shiftKey) {
98+
e.preventDefault();
99+
handleSubmit(e);
100+
}
101+
}}
102+
/>
103+
<div className="flex gap-2 mt-1.5">
104+
<button
105+
type="submit"
106+
disabled={!text.trim() || text.trim() === initialText}
107+
className="rounded px-2.5 py-1 text-[11px] font-medium bg-zinc-700 text-zinc-200 hover:bg-zinc-600 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
108+
>
109+
Save &amp; Submit
110+
</button>
111+
<button
112+
type="button"
113+
onClick={onCancel}
114+
className="rounded px-2.5 py-1 text-[11px] text-zinc-500 hover:text-zinc-300 transition-colors"
115+
>
116+
Cancel
117+
</button>
118+
</div>
119+
</form>
120+
);
121+
}
122+
123+
export function MessageBubble({ message, headers, onRegenerate, onEdit }: MessageBubbleProps) {
59124
const isUser = message.role === 'user';
125+
const [isEditing, setIsEditing] = useState(false);
60126

61127
const role = headers?.['x-ably-role'] ?? message.role;
62128
const clientId = headers?.['x-ably-turn-client-id'];
63129
const turnId = headers?.['x-ably-turn-id'];
64130
const status = headers?.['x-ably-status'];
65131

132+
const messageText = message.parts
133+
.filter((p): p is { type: 'text'; text: string } => p.type === 'text')
134+
.map((p) => p.text)
135+
.join('');
136+
66137
return (
67138
<div className={`flex ${isUser ? 'justify-end' : 'justify-start'}`}>
68139
<div className="max-w-[75%]">
69-
<div className={bubbleClasses(isUser, status)}>
70-
{message.parts.map((part, i) => {
71-
if (part.type === 'text') return <span key={i}>{part.text}</span>;
72-
if (part.type === 'dynamic-tool')
73-
return (
74-
<ToolInvocation
75-
key={i}
76-
part={part as DynamicToolUIPart}
77-
/>
78-
);
79-
return null;
80-
})}
81-
{!isUser && status === 'streaming' && (
82-
<span className="inline-block w-1.5 h-3.5 ml-0.5 bg-amber-500/60 animate-pulse rounded-sm align-text-bottom" />
83-
)}
84-
</div>
85-
<div className="mt-1 flex items-center gap-1.5 flex-wrap">
86-
{/* Regenerate button (assistant messages) */}
87-
{onRegenerate && status !== 'streaming' && (
88-
<button
89-
onClick={onRegenerate}
90-
className="text-[10px] text-zinc-500 hover:text-zinc-200 transition-colors rounded bg-zinc-800/60 px-1.5 py-0.5"
91-
title="Regenerate response"
92-
>
93-
regenerate
94-
</button>
95-
)}
96-
97-
{/* Debug badges */}
98-
{headers && (
99-
<>
100-
<Badge
101-
label="role"
102-
value={role}
103-
color="bg-zinc-900 text-zinc-500"
104-
/>
105-
{clientId && (
106-
<Badge
107-
label="client"
108-
value={clientId}
109-
color="bg-zinc-900 text-zinc-500"
110-
/>
140+
{isEditing && onEdit ? (
141+
<EditForm
142+
initialText={messageText}
143+
onSubmit={(text) => onEdit(text)}
144+
onCancel={() => setIsEditing(false)}
145+
/>
146+
) : (
147+
<>
148+
<div className={bubbleClasses(isUser, status)}>
149+
{message.parts.map((part, i) => {
150+
if (part.type === 'text') return <span key={i}>{part.text}</span>;
151+
if (part.type === 'dynamic-tool')
152+
return (
153+
<ToolInvocation
154+
key={i}
155+
part={part as DynamicToolUIPart}
156+
/>
157+
);
158+
return null;
159+
})}
160+
{!isUser && status === 'streaming' && (
161+
<span className="inline-block w-1.5 h-3.5 ml-0.5 bg-amber-500/60 animate-pulse rounded-sm align-text-bottom" />
162+
)}
163+
</div>
164+
<div className="mt-1 flex items-center gap-1.5 flex-wrap">
165+
{/* Edit button (user messages) */}
166+
{onEdit && status !== 'streaming' && (
167+
<button
168+
onClick={() => setIsEditing(true)}
169+
className="text-[10px] text-zinc-500 hover:text-zinc-200 transition-colors rounded bg-zinc-800/60 px-1.5 py-0.5"
170+
title="Edit message"
171+
>
172+
edit
173+
</button>
111174
)}
112-
{turnId && (
113-
<Badge
114-
label="turn"
115-
value={turnId.slice(0, 8)}
116-
color="bg-zinc-900 text-zinc-500"
117-
/>
175+
176+
{/* Regenerate button (assistant messages) */}
177+
{onRegenerate && status !== 'streaming' && (
178+
<button
179+
onClick={onRegenerate}
180+
className="text-[10px] text-zinc-500 hover:text-zinc-200 transition-colors rounded bg-zinc-800/60 px-1.5 py-0.5"
181+
title="Regenerate response"
182+
>
183+
regenerate
184+
</button>
185+
)}
186+
187+
{/* Debug badges */}
188+
{headers && (
189+
<>
190+
<Badge
191+
label="role"
192+
value={role}
193+
color="bg-zinc-900 text-zinc-500"
194+
/>
195+
{clientId && (
196+
<Badge
197+
label="client"
198+
value={clientId}
199+
color="bg-zinc-900 text-zinc-500"
200+
/>
201+
)}
202+
{turnId && (
203+
<Badge
204+
label="turn"
205+
value={turnId.slice(0, 8)}
206+
color="bg-zinc-900 text-zinc-500"
207+
/>
208+
)}
209+
{status && <StatusBadge status={status} />}
210+
</>
118211
)}
119-
{status && <StatusBadge status={status} />}
120-
</>
121-
)}
122-
</div>
212+
</div>
213+
</>
214+
)}
123215
</div>
124216
</div>
125217
);

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,10 @@ interface MessageListProps {
1111
loading: boolean;
1212
onLoadOlder: () => void;
1313
onRegenerate: (messageId: string) => void;
14+
onEdit: (messageId: string, newText: string) => void;
1415
}
1516

16-
export function MessageList({ nodes, hasOlder, loading, onLoadOlder, onRegenerate }: MessageListProps) {
17+
export function MessageList({ nodes, hasOlder, loading, onLoadOlder, onRegenerate, onEdit }: MessageListProps) {
1718
const endRef = useRef<HTMLDivElement>(null);
1819
const scrollRef = useRef<HTMLDivElement>(null);
1920
const prevLastIdRef = useRef<string | undefined>(undefined);
@@ -61,6 +62,7 @@ export function MessageList({ nodes, hasOlder, loading, onLoadOlder, onRegenerat
6162
message={message}
6263
headers={headers}
6364
onRegenerate={message.role === 'assistant' ? () => onRegenerate(message.id) : undefined}
65+
onEdit={message.role === 'user' ? (text) => onEdit(message.id, text) : undefined}
6466
/>
6567
))}
6668
<div ref={endRef} />

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

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,18 @@
77
* to the core transport's send/cancel methods.
88
*
99
* useChat manages message state before calling sendMessages:
10-
* - submit-message: appends the new user message, passes the full array
10+
* - submit-message (new): appends the new user message, passes the full array
11+
* - submit-message (edit): truncates after the edited message, replaces it,
12+
* passes the truncated array with messageId set
1113
* - regenerate-message: truncates after the target, passes the truncated array
1214
*
1315
* The adapter uses `trigger` to determine the history/messages split:
1416
* - submit-message: last message is new (publish to channel), rest is history
1517
* - regenerate-message: no new messages, entire array is history
18+
*
19+
* When messageId is set (edit or regeneration), the adapter computes fork
20+
* metadata (forkOf/parent) from the conversation tree so the server can
21+
* place the response on the correct branch.
1622
*/
1723

1824
import * as Ably from 'ably';
@@ -35,8 +41,10 @@ export interface SendMessagesRequestContext {
3541
/** What triggered the request: user sent a message, or requested regeneration. */
3642
trigger: 'submit-message' | 'regenerate-message';
3743
/**
38-
* The message ID for regeneration requests. Identifies which assistant
39-
* message to regenerate. Undefined for submit-message.
44+
* The message ID for edit or regeneration requests. For regeneration,
45+
* identifies the assistant message to regenerate. For edits (submit-message
46+
* with messageId), identifies the user message being replaced. Undefined
47+
* when submitting a new message.
4048
*/
4149
messageId?: string;
4250
/** Previous messages in the conversation (context for the LLM). */
@@ -239,11 +247,19 @@ export const createChatTransport = (
239247
}
240248

241249
// Compute fork metadata from the conversation tree.
250+
// For regeneration: messageId is the assistant message being regenerated.
251+
// For edit: messageId is the user message being replaced.
252+
// In both cases: forkOf = the x-ably-msg-id of that message,
253+
// parent = the parent of that message in the tree.
242254
let forkOf: string | undefined;
243255
let parent: string | undefined;
244256

245-
if (trigger === 'regenerate-message' && messageId) {
257+
if (messageId) {
246258
forkOf = messageId;
259+
// Look up the message in the tree to resolve x-ably-msg-id.
260+
// messageId comes from useChat (UIMessage.id) — scan the flattened
261+
// nodes to find the one whose domain message matches this ID.
262+
// Uses the transport's default view — ChatTransport is single-view (one useChat per channel).
247263
const node = transport.view.flattenNodes().find((n) => n.message.id === messageId);
248264
if (node) {
249265
forkOf = node.msgId;

0 commit comments

Comments
 (0)