Stream reattachment client half: resume via same rendering path

Reattach to a turn that outlived its view through the SDK's native
resumeStream() so replay burst + live deltas flow through the exact
same processUIMessageStream chunk pipeline as first-subscriber
streaming — no second rendering dialect.

chat-transport.ts: ReattachingChatTransport subclasses
DefaultChatTransport and repoints reconnectToStream at GET
/api/chats/[id]/stream (the server's reattach endpoint, not the SDK
default /api/chat/{id}/stream). It peeks the first SSE event:
- turn-settled terminal: cancel body, fire onSettled (refetch), return
  null — the SDK treats null as no active stream, status stays ready.
- start chunk whose messageId equals the seeded last assistant message
  id (continuation case): rewritten to a fresh id so the resumed run
  mounts beside the seeded message instead of replacing its parts —
  settle-time record sync merges to the server's assembled message.

chat.tsx: mount effect resumes when initiallyInFlight (no once-ref:
React StrictMode's dev double-mount aborts the first attempt at
dispose — the retry must run). Mid-session discovery sites (jobs
poll, handoff watch) funnel into attachToRunningTurn().
turnInFlight is demoted to fallback liveness: the 1500ms record poll
self-heals a failed attach and clears a stale error banner on settle.
A status watcher clears the flag once the local stream represents
the turn. Stop/Esc unified in abortInFlightTurn: local stop() when a
stream runs here, always POST /api/activity/abort {kind:'turn'} with
409 not_running treated as a benign raced settle. Composer busy now
covers reattached turns (Stop square shows, send gated).

Live-verified on :3001 with real model turns: navigate away mid-turn
and back (replay + live deltas in-view, fallback poll disengaged),
Esc-abort of a reattached turn (partial persisted server-side,
composer restored), two concurrent sessions (raced-settle handled
via turn-settled + record refetch; live settle adopts the final
record), and no leaked stream readers after unmount.
This commit is contained in:
amosblomqvist
2026-09-25 18:12:29 +03:00
parent 7cfd8b5bc7
commit 79e7c8648b
3 changed files with 502 additions and 31 deletions
+199 -2
View File
@@ -21,6 +21,8 @@
* a future "server-ahead" heal in the poller). Rare path; self-heals.
*/
import { DefaultChatTransport } from "ai";
import type { UIMessageChunk } from "ai";
import type { ChatTransport } from "ai";
import type { TeacherUIMessage } from "@/src/core";
/** How often to poll turnInFlight while the racing turn runs. */
@@ -105,8 +107,10 @@ export function retryingFetch(chatId: string, opts: RetryFetchOptions = {}): typ
}
/**
* The chat.tsx one-liner: `transport: createChatTransport(chatId)` replaces
* `new DefaultChatTransport({ api: "/api/chat", body: { chatId } })`.
* The plain transport (no reattach adapters): sends keep the 409
* merge-retry. chat.tsx uses createReattachingTransport below — a
* superset of this — but the plain form stays exported for surfaces
* that only ever send (never resume onto a running turn).
*/
export function createChatTransport(chatId: string): DefaultChatTransport<TeacherUIMessage> {
return new DefaultChatTransport<TeacherUIMessage>({
@@ -114,4 +118,197 @@ export function createChatTransport(chatId: string): DefaultChatTransport<Teache
body: { chatId },
fetch: retryingFetch(chatId),
});
}
/* ────────────────────────────────────────────────────────────────────────
* THE REATTACH TRANSPORT (background-turn runner wave, owner ruling
* 2026-09-25): a turn runs server-side decoupled from every view, and a
* view that arrives late (navigate away mid-turn, wake-turn discovery)
* re-enters it through the SDK's OWN resume path — useChat's
* resumeStream() → reconnectToStream() — so the replay burst and the
* live deltas flow through the same processUIMessageStream assembly as
* first-subscriber streaming. No second rendering dialect exists.
*
* Three adapters over the platform contract, all in this file:
*
* 1. ENDPOINT: the SDK's default reconnect GETs /api/chat/{id}/stream;
* ours is GET /api/chats/[id]/stream (same chunk dialect).
* 2. THE SETTLED TERMINAL: "no turn running" is a terminal EVENT
* (turn-settled + [DONE]), not a 204 — one wire protocol end to end.
* The first SSE event is peeked; turn-settled cancels the body, fires
* onSettled (the Chat clears its in-flight flag and refetches the
* record), and resolves null — the SDK's "no active stream" signal
* (it sets status back to ready, no stream consumed).
* 3. THE CONTINUATION SEAM: the SDK seeds its resume streaming message
* EMPTY (the resume trigger passes lastMessage: undefined), so a
* reattached turn that CONTINUES the seeded last assistant message
* (answer-quiz reaction, wake continuation) would hit the SDK's
* replace-last-message path (start id === last id) and drop that
* message's parts for the whole stream — the answered-quiz card
* would flicker out until the settle sync restores it. When the
* peeked start chunk's messageId equals the seeded last message's id,
* rewrite it to a fresh id: this run's parts mount as their own
* message while the seeded one keeps rendering — pixel-identical to
* the live continuation (card above, reaction below) — and the settle
* syncFromServer lands the server's single assembled message (the
* server is the transcript's canon; the split never persists).
* Part ids stay server-prefixed; nothing keys parts by message id.
* ──────────────────────────────────────────────────────────────────────── */
/** The reattach half of the transport contract (chat.tsx wires the refs). */
export interface ReattachOptions {
/** turn-settled — the stream endpoint found no running turn (raced
* settle, or none ever started): the record is current; the caller
* refetches it and clears its in-flight flag. */
onSettled?: () => void;
/** The LIVE last message id of the seeded transcript (a Chat ref) —
* the continuation detection above. */
lastMessageId?: () => string | undefined;
}
/** Concatenate byte payloads (the head stream's event + read-past remainder). */
function concatBytes(a: Uint8Array, b: Uint8Array): Uint8Array {
const out = new Uint8Array(a.length + b.length);
out.set(a);
out.set(b, a.length);
return out;
}
/** Same id shape the server's runPiTurn mints — collision-free by design. */
function freshMessageId(): string {
return `msg_${Date.now().toString(36)}_${crypto.randomUUID().slice(0, 8)}`;
}
/** One parsed SSE event off the wire (data: <json>\\n\\n framing). */
interface SseEvent {
/** The raw bytes of this event, re-injectable into the head stream. */
raw: Uint8Array;
/** Bytes read PAST the event boundary (the next event can arrive in
* the same read) — they must ride the head of the tail verbatim. */
rest: Uint8Array;
/** The parsed JSON payload (events here are always single data lines). */
data: Record<string, unknown>;
}
/**
* Read the FIRST complete SSE event off a live body (the settled terminal
* and the start chunk are both single small events; a long-lived turn's
* log replays after). Resolves as soon as one \\n\\n boundary is seen — the
* caller re-injects the raw bytes (plus any read-past remainder) ahead of
* the unconsumed tail, so nothing is lost or reordered.
*/
async function readFirstEvent(reader: ReadableStreamDefaultReader<Uint8Array>): Promise<SseEvent> {
const decoder = new TextDecoder();
let buf = "";
for (;;) {
const { done, value } = await reader.read();
if (done) throw new Error("The reattach stream ended before its first event.");
buf += decoder.decode(value, { stream: true });
const end = buf.indexOf("\n\n");
if (end === -1) continue;
const event = buf.slice(0, end);
const dataLine = event.split("\n").find((l) => l.startsWith("data: "));
if (!dataLine) throw new Error("The reattach stream's first event is malformed.");
return {
raw: new TextEncoder().encode(`${event}\n\n`),
rest: new TextEncoder().encode(buf.slice(end + 2)),
data: JSON.parse(dataLine.slice("data: ".length)) as Record<string, unknown>,
};
}
}
/** The reattach transport: DefaultChatTransport for sends, our seam for resumes. */
class ReattachingChatTransport extends DefaultChatTransport<TeacherUIMessage> {
constructor(
private readonly reattach: ReattachOptions,
init: ConstructorParameters<typeof DefaultChatTransport<TeacherUIMessage>>[0],
) {
super(init);
}
override async reconnectToStream(
options: Parameters<ChatTransport<TeacherUIMessage>["reconnectToStream"]>[0],
): Promise<ReadableStream<UIMessageChunk> | null> {
const res = await fetch(`/api/chats/${encodeURIComponent(options.chatId)}/stream`, {
signal: options.abortSignal,
});
if (res.status === 204) return null; // the SDK's native no-stream signal
if (!res.ok) throw new Error((await res.text()) || "Failed to attach to the running turn.");
if (!res.body) throw new Error("The reattach response body is empty.");
const reader = res.body.getReader();
let first: SseEvent;
try {
first = await readFirstEvent(reader);
} catch (err) {
void reader.cancel().catch(() => {});
throw err;
}
// The settled terminal: nothing to attach to — the record is current.
if (first.data.type === "turn-settled") {
await reader.cancel().catch(() => {});
this.reattach.onSettled?.();
return null;
}
// The continuation seam (see the block comment): a start chunk whose
// id matches the seeded last message would REPLACE that message with
// an empty streaming state — rewrite to a fresh id so this run mounts
// beside it and the settle sync merges the canon back.
let firstEvent = first.raw;
if (
first.data.type === "start" &&
typeof first.data.messageId === "string" &&
first.data.messageId === this.reattach.lastMessageId?.()
) {
firstEvent = new TextEncoder().encode(
`data: ${JSON.stringify({ ...first.data, messageId: freshMessageId() })}\n\n`,
);
}
// The head bytes: the (possibly rewritten) first event plus anything
// the peek read past the boundary — then the live tail, untouched.
const head = first.rest.length > 0 ? concatBytes(firstEvent, first.rest) : firstEvent;
// Re-inject the (possibly rewritten) head ahead of the live tail and
// hand the SDK the SAME SSE parse the POST response goes through.
const tail = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(head);
void (async () => {
try {
for (;;) {
const { done, value } = await reader.read();
if (done) {
controller.close();
return;
}
controller.enqueue(value);
}
} catch (err) {
controller.error(err);
}
})();
},
cancel() {
return reader.cancel();
},
});
return this.processResponseStream(tail);
}
}
/**
* The chat.tsx wiring: `transport: createReattachingTransport(chatId, refs)`.
* Sends keep the 409 merge-retry fetch; resumes take the adapters above.
*/
export function createReattachingTransport(
chatId: string,
reattach: ReattachOptions,
): DefaultChatTransport<TeacherUIMessage> {
return new ReattachingChatTransport(reattach, {
api: "/api/chat",
body: { chatId },
fetch: retryingFetch(chatId),
});
}