Continuation replay duplicates already-rendered assistant parts after reconnect

[cgrdavies]

Describe the bug

A WebSocket reconnect during an active continuation stream can duplicate content that the client already rendered.

This happens with a single legitimate resume handshake and full-buffer replay; it does not require the duplicate-STREAM_RESUMING / double-ACK condition fixed by #1742. The duplicate content is appended inside the existing assistant message and may disappear later when an authoritative message update replaces the temporary client state.

Relationship to #1733 / #1742

This is adjacent to, but distinct from #1733:

• ai-chat: double STREAM_RESUMING replay duplicates the streaming assistant text part #1733 mechanism: two STREAM_RESUMING notifications caused two ACKs and two complete buffer replays; the second replay stacked a duplicate text part.
• This mechanism: two notifications may still arrive, but the client sends only one ACK and receives one complete replay. That single replay overlaps continuation content already rendered live before reconnect.

The captured frame sequence for this case contains exactly one STREAM_RESUME_ACK and one replay batch, so #1742's ACK deduplication is working. #1742 makes non-continuation replay idempotent, but deliberately excludes continuation replay from the whole-message reset to preserve completed pre-continuation steps. This report covers that remaining continuation-specific gap.

Both React paths appear affected:

• Transport-owned resumes snapshot the existing trailing assistant and feed the replayed continuation chunks into AI SDK useChat.
• The fallback observer initializes its continuation accumulator from the assistant's current parts, then applies the complete replay buffer.

To Reproduce

1. Render chat with useAgentChat and resumable streaming enabled.
2. Start a turn that enters a continuation after at least one completed step/tool call.
3. Allow part of the active continuation to render.
4. Reconnect the socket while that continuation is still streaming. With an async useAgent query, the default five-minute cache refresh is one deterministic way to trigger this; forcing a reconnect also works.
5. Complete one STREAM_RESUMING -> STREAM_RESUME_ACK handshake and replay the stored continuation buffer.

Minimal state-machine shape:

Current assistant parts:
  "before continuation"
  "already streamed"

Replayed frames:
  start          (replay=true, continuation=true, no messageId)
  text-start
  text-delta     "already streamed"
  replayComplete

Actual assistant parts:
  "before continuation"
  "already streamed"
  "already streamed"

The same result reproduces through AI SDK React's Chat.resumeStream() when the current assistant already contains the live continuation suffix.

Expected behavior

Stream replay should be idempotent. The continuation suffix already rendered before reconnect should appear once, while all assistant parts completed before the continuation remain intact.

Screenshots

Not included because the reducer-level reproduction is deterministic, but a recording can be provided if useful.

Version:

• agents@0.17.3 and agents@0.17.4
• @cloudflare/ai-chat@0.9.3
• Reproduced with @ai-sdk/react@3.0.219 / ai@6.0.217

Additional context

The protocol currently creates a gap between message identity and replay identity:

• Continuation start chunks intentionally have messageId removed before storage/broadcast:

  agents/packages/ai-chat/src/index.ts

  Lines 6227 to 6284 in 03cdc82

  	// Rewrite chunks before storing and broadcasting:
  	// 1. Strip messageId from continuation start chunks so clients
  	// reuse the existing assistant message (#1229).
  	// 2. Stamp the allocated assistant id onto a new turn's start chunk
  	// so the client builds the live message under the SAME id the
  	// server persists under (see below).
  	// 3. Convert the internal "finish" event's finishReason into the
  	// UIMessageStreamPart messageMetadata format (#677).
  	let eventToSend: unknown = data;
  	if (data.type === "start") {
  	if (continuation && "messageId" in data) {
  	const { messageId: _, ...rest } = data as {
  	messageId: unknown;
  	[key: string]: unknown;
  	};
  	eventToSend = rest;
  	} else if (!continuation) {
  	// Most providers (e.g. Workers AI) emit no `start.messageId`,
  	// so the client's AI SDK would build the streaming assistant
  	// under its own generated id while the server persists under
  	// `message.id`. The two then can't be reconciled by id, and the
  	// originating tab briefly renders the turn twice — the live copy
  	// plus the `CF_AGENT_CHAT_MESSAGES` broadcast — before
  	// collapsing. Stamping the allocated id here makes the common
  	// case behave like the provider-id case the client already
  	// relies on (react.tsx records `start.messageId` to map the
  	// local stream to the persisted message).
  	const startData = data as {
  	messageId?: unknown;
  	[key: string]: unknown;
  	};
  	if (startData.messageId == null) {
  	eventToSend = { ...startData, messageId: message.id };
  	}
  	}
  	}
  	if (data.type === "finish" && "finishReason" in data) {
  	const { finishReason, ...rest } = data as {
  	finishReason: string;
  	[key: string]: unknown;
  	};
  	eventToSend = {
  	...rest,
  	type: "finish",
  	messageMetadata: { finishReason }
  	};
  	}
  	// Store chunk for replay and broadcast to clients
  	const chunkBody = JSON.stringify(eventToSend);
  	await this._storeStreamChunk(streamId, chunkBody);
  	this._broadcastChatMessage({
  	body: chunkBody,
  	done: false,
  	id,
  	type: MessageType.CF_AGENT_USE_CHAT_RESPONSE,
  	...(continuation && { continuation: true })
  	});

• The React replay reset requires a matching start.messageId and explicitly excludes continuation replays so it does not erase pre-continuation parts:

  agents/packages/agents/src/chat/react.tsx

  Lines 2025 to 2082 in 03cdc82

  	case MessageType.CF_AGENT_USE_CHAT_RESPONSE: {
  	if (localRequestIdsRef.current.has(data.id)) {
  	if (data.body?.trim()) {
  	try {
  	const chunkData = JSON.parse(data.body) as {
  	messageId?: string;
  	type?: string;
  	};
  	if (
  	chunkData.type === "start" &&
  	typeof chunkData.messageId === "string"
  	) {
  	localResponseIds.set(data.id, chunkData.messageId);
  	// Re-arm streaming protection to the ACTUAL assistant id for
  	// this turn. `protectStreamingAssistantTail` runs at send
  	// time — before the assistant message is minted — so it can
  	// only latch the PREVIOUS turn's id (or nothing on the first
  	// turn). Without correcting it here, a mid-stream full-list
  	// broadcast (`CF_AGENT_CHAT_MESSAGES`, which Think emits after
  	// every tool result) replaces the live-streamed assistant
  	// with a possibly-behind server snapshot, so its parts (e.g.
  	// tool cards) briefly disappear and reappear. Continuations
  	// are skipped: they extend the existing protected assistant.
  	if (!data.continuation) {
  	const protection = protectedStreamingAssistantRef.current;
  	if (protection?.assistantId !== chunkData.messageId) {
  	const msgs = messagesRef.current;
  	const idx = msgs.findIndex(
  	(m) => m.id === chunkData.messageId
  	);
  	const anchorMessageId =
  	idx >= 0
  	? (msgs[idx - 1]?.id ?? null)
  	: (msgs[msgs.length - 1]?.id ?? null);
  	protectedStreamingAssistantRef.current = {
  	assistantId: chunkData.messageId,
  	anchorMessageId
  	};
  	}
  	}
  	// EVERY replayed `start` rebuilds the message from chunk 0,
  	// so the matching trailing assistant must be reset each
  	// time — not only while the resume request id is still
  	// pending (#1733: a second replay otherwise stacks a
  	// duplicate text part). Continuation replays are excluded:
  	// they append to the existing assistant message, and
  	// wiping it would drop the pre-continuation parts.
  	if (
  	data.replay &&
  	!data.continuation &&
  	!resumingToolContinuationRef.current &&
  	observedToolContinuationRequestIdRef.current !== data.id
  	) {
  	pendingReplayResumeRequestIdsRef.current.delete(data.id);
  	resetMatchingHydratedAssistantForReplay(
  	chunkData.messageId
  	);
  	}

• The fallback accumulator seeds continuation replay from the assistant's current parts, which can already include live bytes from this same continuation:

  agents/packages/agents/src/chat/broadcast-state.ts

  Lines 84 to 139 in 03cdc82

  	case "response": {
  	let accumulator: StreamAccumulator;
  	// A replayed `start` chunk means the server is re-sending the stream
  	// buffer from chunk 0 (resume replay). Re-initialize the accumulator
  	// instead of appending into an existing one: replaying into an
  	// accumulator that already holds this stream's parts would duplicate
  	// them (a second `text-start` unconditionally opens a second text
  	// part — #1733). Re-initializing makes replay idempotent under any
  	// number of replays, including a second replay triggered by a
  	// duplicate STREAM_RESUMING → ACK cycle or a reconnect.
  	const isReplayedStart =
  	event.replay === true &&
  	(event.chunkData as { type?: string } | null | undefined)?.type ===
  	"start";
  	if (
  	state.status === "idle" ||
  	state.streamId !== event.streamId ||
  	isReplayedStart
  	) {
  	let messageId = event.messageId;
  	let existingParts: UIMessage["parts"] | undefined;
  	let existingMetadata: Record<string, unknown> | undefined;
  	if (event.continuation && event.currentMessages) {
  	for (let i = event.currentMessages.length - 1; i >= 0; i--) {
  	if (event.currentMessages[i].role === "assistant") {
  	messageId = event.currentMessages[i].id;
  	existingParts = [...event.currentMessages[i].parts];
  	if (event.currentMessages[i].metadata != null) {
  	existingMetadata = {
  	...(event.currentMessages[i].metadata as Record<
  	string,
  	unknown
  	>)
  	};
  	}
  	break;
  	}
  	}
  	}
  	accumulator = new StreamAccumulator({
  	messageId,
  	continuation: event.continuation,
  	existingParts,
  	existingMetadata
  	});
  	} else {
  	accumulator = state.accumulator;
  	}
  	if (event.chunkData) {
  	accumulator.applyChunk(event.chunkData as StreamChunkData);

The existing continuation replay regression test seeds only pre-continuation content. It does not include a continuation suffix that was rendered live before reconnect, so it misses the overlap:

agents/packages/agents/src/chat/__tests__/broadcast-state.test.ts

Lines 552 to 598 in 03cdc82

	it("replayed continuation start re-initializes with existing parts intact", () => {
	const current = makeMessages("hi", "first half");
	const replay = (state: BroadcastStreamState, chunkData: unknown) =>
	transition(state, {
	type: "response",
	streamId: "req-c",
	messageId: "tmp",
	chunkData,
	replay: true,
	continuation: true,
	currentMessages: current
	});
	let state = replay(idle, { type: "start", messageId: "msg-1" }).state;
	state = replay(state, {
	type: "text-delta",
	id: "t1",
	delta: " second half"
	}).state;
	expect(state.status).toBe("observing");
	if (state.status === "observing") {
	// Continuation picked up the trailing assistant's id + parts.
	expect(state.accumulator.messageId).toBe("msg-1");
	const texts = state.accumulator.parts.filter((p) => p.type === "text");
	expect(texts.length).toBeGreaterThan(0);
	}
	const done = transition(state, {
	type: "response",
	streamId: "req-c",
	messageId: "tmp",
	done: true,
	continuation: true,
	currentMessages: current
	});
	const messages = done.messagesUpdate!(current);
	// Continuation merged into the existing assistant, no extra message.
	expect(messages).toHaveLength(2);
	const fullText = messages[1].parts
	.filter((p) => p.type === "text")
	.map((p) => (p as { text: string }).text)
	.join("");
	expect(fullText).toContain("first half");
	expect(fullText).toContain(" second half");
	});

A possible fix is to preserve a per-request pre-continuation baseline and restore it before applying the complete replay, or introduce an explicit segment/chunk identity so replay can replace or deduplicate only the current continuation suffix. Clearing the entire assistant message would lose completed pre-continuation steps.

Regression coverage should exercise both the transport-owned and fallback paths with:

pre-continuation parts + already-rendered continuation suffix
-> reconnect
-> one ACK
-> continuation replay starting without messageId
-> suffix appears exactly once