From d45d7a4728df927ef20775a40fc5eb2b58f085cb Mon Sep 17 00:00:00 2001 From: Robel Estifanos Date: Sun, 23 Aug 2026 18:34:08 -0400 Subject: [PATCH 1/2] fix(streaming): drain deltas at end-of-stream, not at the final step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit saveStreamDeltas discarded buffered deltas: the final step's onStepEnd marked the stream finished externally, which stopped addParts and skipped finish() — the only thing that drains the buffer. A short generation could persist a single `start` part while the saved message was complete. The deeper problem is that onStepEnd cannot see the data. Parts live inside the AI SDK's transform and tee pipeline until consumeStream's iterator yields them, so no callback fires with everything handed over — the stream-level `finish` chunk provably arrives after the last onStepEnd. Split the one flag into the two things it conflated. Finish ownership is decided at construction; "stop accepting parts" applies only where the row is finished before the source is drained. consumeStream now flushes at EOF, which is the only point where every part has actually been handed over, and the awaited path leaves the row streaming until then — so it now persists the trailing chunks it used to drop. returnImmediately keeps flushing at onStepEnd, since nothing awaits consumption there and #265 requires the inline save. Parts still inside the SDK at that instant are not persisted as deltas; the atomically saved message is authoritative. Fixes #323. --- src/vercel/client/streamText.test.ts | 139 ++++++++++++++++++ src/vercel/client/streamText.ts | 21 ++- .../client/streaming.integration.test.ts | 43 +++++- src/vercel/client/streaming.test.ts | 30 ++++ src/vercel/client/streaming.ts | 112 +++++++++++--- 5 files changed, 320 insertions(+), 25 deletions(-) diff --git a/src/vercel/client/streamText.test.ts b/src/vercel/client/streamText.test.ts index 4328011b..b078b205 100644 --- a/src/vercel/client/streamText.test.ts +++ b/src/vercel/client/streamText.test.ts @@ -91,9 +91,51 @@ export const streamTextEmptyReturnImmediately = action({ }, }); +export const streamTextThrottled = action({ + args: { threadId: v.string() }, + handler: async (ctx, { threadId }) => { + const result = await agent.streamText( + ctx, + { threadId }, + { prompt: "Test" }, + { + saveStreamDeltas: { + returnImmediately: true, + chunking: "word", + throttleMs: 60_000, + }, + }, + ); + await result.consumeStream(); + return { ok: true }; + }, +}); + +// Same as streamTextThrottled, but awaited: streamText consumes the stream +// itself, so the terminal transition happens at end-of-stream. +export const streamTextThrottledAwaited = action({ + args: { threadId: v.string() }, + handler: async (ctx, { threadId }) => { + await agent.streamText( + ctx, + { threadId }, + { prompt: "Test" }, + { + saveStreamDeltas: { + chunking: "word", + throttleMs: 60_000, + }, + }, + ); + return { ok: true }; + }, +}); + const testApi: ApiFromModules<{ fns: { streamTextReturnImmediately: typeof streamTextReturnImmediately; + streamTextThrottled: typeof streamTextThrottled; + streamTextThrottledAwaited: typeof streamTextThrottledAwaited; streamTextEmptyAwaited: typeof streamTextEmptyAwaited; streamTextEmptyReturnImmediately: typeof streamTextEmptyReturnImmediately; }; @@ -221,3 +263,100 @@ describe("streamText with an empty final step (issue #274)", () => { }, ); }); + +describe("saveStreamDeltas flushes buffered parts (issue #323)", () => { + test("deltas hold the full text when the generation outpaces the throttle", async () => { + const t = initConvexTest(schema); + const threadId = await t.run(async (ctx) => + createThread(ctx, components.agent, { userId: "u1" }), + ); + + await t.action(testApi.streamTextThrottled, { threadId }); + await t.finishAllScheduledFunctions(() => {}); + + const streams = await t.run(async (ctx) => + ctx.runQuery(components.agent.streams.list, { + threadId, + statuses: ["streaming", "finished", "aborted"], + }), + ); + const deltas = await t.run(async (ctx) => + ctx.runQuery(components.agent.streams.listDeltas, { + threadId, + cursors: streams.map((s) => ({ streamId: s.streamId, cursor: 0 })), + }), + ); + + expect(streams).toHaveLength(1); + expect(streams[0].status).toBe("finished"); + + let cursor = 0; + for (const delta of deltas) { + expect(delta.start).toBe(cursor); + cursor = delta.end; + } + + const parts = deltas.flatMap((d) => d.parts); + const types = parts.map((p) => p.type); + expect(types.at(0)).toBe("start"); + expect(types).toContain("text-start"); + expect(types).toContain("text-end"); + // The stream-level "finish" chunk is emitted after the last step ends, so + // it cannot exist yet; the row's finished status carries that instead. + expect(types.at(-1)).toBe("finish-step"); + expect(types).not.toContain("finish"); + expect( + parts + .filter((p) => p.type === "text-delta") + .map((p) => (p as { delta?: string }).delta ?? "") + .join(""), + ).toBe(FINAL_TEXT); + }); + + test("the awaited path captures the stream-level finish chunk", async () => { + const t = initConvexTest(schema); + const threadId = await t.run(async (ctx) => + createThread(ctx, components.agent, { userId: "u1" }), + ); + + await t.action(testApi.streamTextThrottledAwaited, { threadId }); + await t.finishAllScheduledFunctions(() => {}); + + const streams = await t.run(async (ctx) => + ctx.runQuery(components.agent.streams.list, { + threadId, + statuses: ["streaming", "finished", "aborted"], + }), + ); + const deltas = await t.run(async (ctx) => + ctx.runQuery(components.agent.streams.listDeltas, { + threadId, + cursors: streams.map((s) => ({ streamId: s.streamId, cursor: 0 })), + }), + ); + + expect(streams).toHaveLength(1); + expect(streams[0].status).toBe("finished"); + + let cursor = 0; + for (const delta of deltas) { + expect(delta.start).toBe(cursor); + cursor = delta.end; + } + + const parts = deltas.flatMap((d) => d.parts); + const types = parts.map((p) => p.type); + // Unlike the returnImmediately path, nothing stops accepting parts early + // here: consumeStream drains at EOF, so the trailing chunks the AI SDK + // emits after the last onStepEnd are persisted too. + expect(types.at(0)).toBe("start"); + expect(types.at(-1)).toBe("finish"); + expect(types).toContain("finish-step"); + expect( + parts + .filter((p) => p.type === "text-delta") + .map((p) => (p as { delta?: string }).delta ?? "") + .join(""), + ).toBe(FINAL_TEXT); + }); +}); diff --git a/src/vercel/client/streamText.ts b/src/vercel/client/streamText.ts index af2bc15c..3b5b11db 100644 --- a/src/vercel/client/streamText.ts +++ b/src/vercel/client/streamText.ts @@ -156,6 +156,10 @@ export async function streamText< materialize: (parts) => materializeUIMessageChunkFiles(ctx, component, parts), abortSignal: args.abortSignal, + // Both paths below finish the stream row atomically with the + // message save (issue #181), so the streamer must never finish it + // itself at end-of-stream. + finishHandledExternally: true, }, { threadId, @@ -215,10 +219,13 @@ export async function streamText< const createPendingMessage = await willContinue(steps, args.stopWhen); if (!createPendingMessage && streamer) { // Final step with streaming enabled. - streamer.markFinishedExternally(); if (willAwaitStream) { // We're about to `await stream` below — defer the save so it - // happens atomically with stream finish (issue #181). + // happens atomically with stream finish (issue #181). Don't touch + // the streamer here: the stream-level `finish` chunk is emitted + // after this callback, so the row has to stay `streaming` and keep + // accepting parts until consumeStream reaches EOF, which is the only + // point where everything has actually been handed over. pendingFinalStep = { step, responseMessages: responseMessagesForStep(step), @@ -226,7 +233,10 @@ export async function streamText< } else { // returnImmediately path: streamText is about to return without // awaiting consumption, so the deferred-save block below won't - // see this step. Save inline now (issue #265). + // see this step. Save inline now (issue #265). Nothing awaits the + // stream here, so this is the last moment we can drain deltas — see + // flushAndStopAccepting for the window that leaves. + await streamer.flushAndStopAccepting(); const finishStreamId = await streamer.getOrCreateStreamId(); await call.save( { step, responseMessages: responseMessagesForStep(step) }, @@ -276,6 +286,11 @@ export async function streamText< if (pendingFinalStep && streamer) { const finishStreamId = await streamer.getOrCreateStreamId(); await call.save(pendingFinalStep, false, finishStreamId); + } else if (willAwaitStream && streamer) { + // No final step was deferred (e.g. the generation produced none), so no + // save will finish the stream. The streamer doesn't finish itself, so do + // it here rather than leaving the row to time out. + await streamer.finish(); } const metadata: GenerationOutputMetadata = { promptMessageId, diff --git a/src/vercel/client/streaming.integration.test.ts b/src/vercel/client/streaming.integration.test.ts index bee17e31..8b1a01ea 100644 --- a/src/vercel/client/streaming.integration.test.ts +++ b/src/vercel/client/streaming.integration.test.ts @@ -190,17 +190,16 @@ describe("HTTP Streaming Initiation", () => { }); }); - test("markFinishedExternally prevents consumeStream from calling finish", async () => { + test("finishHandledExternally prevents consumeStream from calling finish", async () => { await t.run(async (ctx) => { const streamer = new DeltaStreamer( components.agent, ctx, - { ...defaultTestOptions }, + { ...defaultTestOptions, finishHandledExternally: true }, { ...testMetadata, threadId }, ); await streamer.getStreamId(); - streamer.markFinishedExternally(); const result = streamText({ model: mockModel({ @@ -217,6 +216,44 @@ describe("HTTP Streaming Initiation", () => { { threadId, statuses: ["streaming"] }, ); expect(streamingStreams).toHaveLength(1); + + // ...but the parts still made it into deltas, since the row kept + // accepting until end-of-stream. + const deltas = await ctx.runQuery(components.agent.streams.listDeltas, { + threadId, + cursors: [{ streamId: streamer.streamId!, cursor: 0 }], + }); + const types = deltas.flatMap((d) => d.parts).map((p) => p.type); + expect(types.at(0)).toBe("start"); + expect(types.at(-1)).toBe("finish"); + }); + }); + + test("flushAndStopAccepting drops later parts but keeps the row streaming", async () => { + await t.run(async (ctx) => { + const streamer = new DeltaStreamer( + components.agent, + ctx, + { ...defaultTestOptions, finishHandledExternally: true }, + { ...testMetadata, threadId }, + ); + + await streamer.addParts([{ type: "start" }]); + await streamer.flushAndStopAccepting(); + await streamer.addParts([{ type: "finish" }]); + + const deltas = await ctx.runQuery(components.agent.streams.listDeltas, { + threadId, + cursors: [{ streamId: streamer.streamId!, cursor: 0 }], + }); + const types = deltas.flatMap((d) => d.parts).map((p) => p.type); + expect(types).toEqual(["start"]); + + const streamingStreams = await ctx.runQuery( + components.agent.streams.list, + { threadId, statuses: ["streaming"] }, + ); + expect(streamingStreams).toHaveLength(1); }); }); }); diff --git a/src/vercel/client/streaming.test.ts b/src/vercel/client/streaming.test.ts index b10baf70..b0b1f9a6 100644 --- a/src/vercel/client/streaming.test.ts +++ b/src/vercel/client/streaming.test.ts @@ -428,4 +428,34 @@ describe("DeltaStreamer", () => { ); }); // TODO: test fetching partial stream data - syncStreams w/ cursors + + test("does not drop a part parked in stream creation when the final step lands", async () => { + let resolveCreate!: (streamId: string) => void; + const creating = new Promise((r) => (resolveCreate = r)); + const sent: unknown[] = []; + const runMutation = vi + .fn() + .mockImplementationOnce(() => creating) + .mockImplementation((_ref: unknown, args: unknown) => { + sent.push(args); + return Promise.resolve(true); + }); + const streamer = new DeltaStreamer( + components.agent, + { runMutation } as unknown as MutationCtx, + { ...defaultTestOptions }, + { ...testMetadata, threadId }, + ); + + // A part arrives and parks in streams.create. + const adding = streamer.addParts(["A"]); + // The final step lands while creation is still in flight. + const finishing = streamer.flushAndStopAccepting(); + resolveCreate("stream-1"); + await Promise.all([adding, finishing]); + + expect(sent).toHaveLength(1); + expect((sent[0] as { parts: string[] }).parts).toEqual(["A"]); + }); + }); diff --git a/src/vercel/client/streaming.ts b/src/vercel/client/streaming.ts index 5f9961c3..27334beb 100644 --- a/src/vercel/client/streaming.ts +++ b/src/vercel/client/streaming.ts @@ -217,9 +217,21 @@ export class DeltaStreamer { #abortPromise: Promise | undefined; #cursor: number = 0; public abortController: AbortController; - // When true, the stream will be finished externally (e.g., atomically via addMessages) - // and consumeStream should skip calling finish(). - #finishedExternally: boolean = false; + /** + * Who owns the terminal `streams.finish` transition. When true, external code + * finishes the stream row (e.g. atomically with the message save in + * `addMessages`, see issue #181), so `consumeStream` must not finish it. + * Decided once, at construction, by whoever owns the generation. + */ + #finishHandledExternally: boolean; + /** + * Whether we've stopped accepting new parts. Only meaningful when the stream + * row is finished before the source stream has been fully consumed (the + * `returnImmediately` path). This is an optimization, not a race guard: + * `streams.addDelta` already returns false for a non-streaming row and + * `#sendDelta` treats that as a benign late miss. + */ + #stoppedAccepting: boolean = false; constructor( public readonly component: AgentComponent, @@ -229,6 +241,13 @@ export class DeltaStreamer { onAsyncAbort: (reason: string) => Promise; abortSignal: AbortSignal | undefined; compress: ((parts: T[]) => T[]) | null; + /** + * Set when external code will call `streams.finish` (typically + * atomically with the message save). `consumeStream` then flushes at EOF + * but leaves the terminal transition to that caller. Defaults to false: + * `consumeStream` finishes the stream itself. + */ + finishHandledExternally?: boolean; materialize?: (parts: T[]) => Promise<{ parts: T[]; fileRefs: Array<{ url: string; fileId: string }>; @@ -252,6 +271,7 @@ export class DeltaStreamer { compress: config.compress, materialize: config.materialize ?? null, }; + this.#finishHandledExternally = config.finishHandledExternally ?? false; this.#nextParts = []; this.abortController = new AbortController(); if (config.abortSignal) { @@ -293,15 +313,21 @@ export class DeltaStreamer { if (this.abortController.signal.aborted) { return; } - // Once the stream has been finished externally (e.g. by the inline - // save in streamText's onStepFinish for the returnImmediately path), - // the stream record is already "finished" in the DB. Late deltas - // would be silently dropped by streams.addDelta — skip the work. - if (this.#finishedExternally) { + // Once we've stopped accepting (the returnImmediately path finished the + // stream row inline), the stream record is already "finished" in the DB. + // Late deltas would be silently dropped by streams.addDelta — skip the work. + if (this.#stoppedAccepting) { return; } - await this.getStreamId(); + // Buffer before awaiting: a part parked in stream creation would be + // invisible to #flushPendingParts, which only inspects the buffer and the + // in-flight write. The check above and this push are synchronous, so a + // concurrent flush either sees this part or has yet to set the flag. this.#nextParts.push(...parts); + await this.getStreamId(); + if (this.#stoppedAccepting || this.abortController.signal.aborted) { + return; + } if ( !this.#ongoingWrite && Date.now() - this.#latestWrite >= this.config.throttleMs @@ -324,24 +350,71 @@ export class DeltaStreamer { ).catch(() => {}); throw error; } - // Skip finish if it will be handled externally (atomically with message save) - // or if the stream was aborted (e.g., due to a failed delta write). // Abort cleanup owns the terminal component transition, so consumeStream // must wait for it instead of also trying to finish the stream. if (this.abortController.signal.aborted) { await this.#waitForAbortCleanup(); - } else if (!this.#finishedExternally) { + return; + } + // EOF is the only point at which every part the source will ever hand over + // has actually been handed over — parts live inside the AI SDK's transform + // and tee pipeline until the iterator yields them, so no earlier callback + // can observe them. Drain here, before the terminal transition. + try { + await this.#flushPendingParts(); + } catch (error) { + if (this.abortController.signal.aborted) { + await this.#waitForAbortCleanup(); + } + throw error; + } + if (this.abortController.signal.aborted) { + await this.#waitForAbortCleanup(); + return; + } + // Skip finish if it will be handled externally (atomically with the + // message save). + if (!this.#finishHandledExternally) { await this.finish(); } } /** - * Mark the stream as being finished externally (e.g., atomically via addMessages). - * When called, consumeStream() will skip calling finish() since it will be - * handled elsewhere in the same mutation as message saving. + * Drain everything currently buffered or in flight, so that after this + * resolves no delta write is outstanding and #nextParts is empty. + */ + async #flushPendingParts(): Promise { + while (!this.abortController.signal.aborted) { + const inFlight = this.#ongoingWrite; + await inFlight; + // #sendDelta reassigns #ongoingWrite from its own tail, so a write can + // still be live even though the buffer it drained is now empty. + if (this.#ongoingWrite !== inFlight) { + continue; + } + if (this.#nextParts.length === 0) { + break; + } + this.#ongoingWrite = this.#sendDelta(); + } + } + + /** + * Flush buffered parts and stop accepting new ones, for callers that finish + * the stream row before the source stream has been fully consumed — i.e. the + * `returnImmediately` path, where nothing awaits consumption and the save has + * to happen inline (issue #265). + * + * There is an inherent window here: parts still inside the AI SDK's + * transform/tee pipeline at this instant have not reached addParts, so they + * are never persisted as deltas (the stream-level `finish` chunk among them). + * That is acceptable because the message saved alongside this transition is + * complete and atomic — the saved message, not the delta log, is + * authoritative once the stream is finished. */ - public markFinishedExternally(): void { - this.#finishedExternally = true; + public async flushAndStopAccepting(): Promise { + await this.#flushPendingParts(); + this.#stoppedAccepting = true; } /** @@ -358,6 +431,7 @@ export class DeltaStreamer { } let success: boolean; try { + await this.getStreamId(); const delta = await this.#createDelta(); if (!delta) { return; @@ -372,11 +446,11 @@ export class DeltaStreamer { return; } if (!success) { - // An in-flight #sendDelta started before markFinishedExternally() + // A #sendDelta racing the inline save on the returnImmediately path // will get `success === false` because the stream row is already // "finished". That's a benign late-write miss, not a failure — // don't convert it into an abort. - if (this.#finishedExternally) { + if (this.#stoppedAccepting) { return; } await this.#abortDelta("async abort"); From fb48ae7df72097d5323e674b274ed60f501f2648 Mon Sep 17 00:00:00 2001 From: Robel Estifanos Date: Tue, 25 Aug 2026 15:15:53 -0400 Subject: [PATCH 2/2] fix(streaming): finish the stream row when messages aren't stored MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With storageOptions.saveMessages set to "none", call.save never reaches addMessages, so nothing performed the atomic finish and the row stayed "streaming" until its timeout. Both paths were affected. Derive finish ownership from whether messages are stored rather than assuming a save will always follow, and finish explicitly after the inline flush on the returnImmediately path. Predates this branch — reproduces on main. --- src/vercel/client/streamText.test.ts | 77 ++++++++++++++++++++++++++++ src/vercel/client/streamText.ts | 12 +++-- src/vercel/client/streaming.ts | 54 ++++++------------- 3 files changed, 102 insertions(+), 41 deletions(-) diff --git a/src/vercel/client/streamText.test.ts b/src/vercel/client/streamText.test.ts index b078b205..ce794944 100644 --- a/src/vercel/client/streamText.test.ts +++ b/src/vercel/client/streamText.test.ts @@ -131,11 +131,50 @@ export const streamTextThrottledAwaited = action({ }, }); +export const streamTextNoStorage = action({ + args: { threadId: v.string() }, + handler: async (ctx, { threadId }) => { + await agent.streamText( + ctx, + { threadId }, + { prompt: "Test" }, + { + saveStreamDeltas: { chunking: "word", throttleMs: 0 }, + storageOptions: { saveMessages: "none" }, + }, + ); + return { ok: true }; + }, +}); + +export const streamTextNoStorageImmediate = action({ + args: { threadId: v.string() }, + handler: async (ctx, { threadId }) => { + const r = await agent.streamText( + ctx, + { threadId }, + { prompt: "Test" }, + { + saveStreamDeltas: { + returnImmediately: true, + chunking: "word", + throttleMs: 0, + }, + storageOptions: { saveMessages: "none" }, + }, + ); + await r.consumeStream(); + return { ok: true }; + }, +}); + const testApi: ApiFromModules<{ fns: { streamTextReturnImmediately: typeof streamTextReturnImmediately; streamTextThrottled: typeof streamTextThrottled; streamTextThrottledAwaited: typeof streamTextThrottledAwaited; + streamTextNoStorage: typeof streamTextNoStorage; + streamTextNoStorageImmediate: typeof streamTextNoStorageImmediate; streamTextEmptyAwaited: typeof streamTextEmptyAwaited; streamTextEmptyReturnImmediately: typeof streamTextEmptyReturnImmediately; }; @@ -360,3 +399,41 @@ describe("saveStreamDeltas flushes buffered parts (issue #323)", () => { ).toBe(FINAL_TEXT); }); }); + +describe("stream finish ownership without message storage", () => { + test("the row still terminates when saveMessages is none", async () => { + const t = initConvexTest(schema); + const threadId = await t.run(async (ctx) => + createThread(ctx, components.agent, { userId: "u1" }), + ); + + await t.action(testApi.streamTextNoStorage, { threadId }); + await t.finishAllScheduledFunctions(() => {}); + + const streams = await t.run(async (ctx) => + ctx.runQuery(components.agent.streams.list, { + threadId, + statuses: ["streaming", "finished", "aborted"], + }), + ); + expect(streams.map((s) => s.status)).toEqual(["finished"]); + }); + + test("the row still terminates on the returnImmediately path", async () => { + const t = initConvexTest(schema); + const threadId = await t.run(async (ctx) => + createThread(ctx, components.agent, { userId: "u1" }), + ); + + await t.action(testApi.streamTextNoStorageImmediate, { threadId }); + await t.finishAllScheduledFunctions(() => {}); + + const streams = await t.run(async (ctx) => + ctx.runQuery(components.agent.streams.list, { + threadId, + statuses: ["streaming", "finished", "aborted"], + }), + ); + expect(streams.map((s) => s.status)).toEqual(["finished"]); + }); +}); diff --git a/src/vercel/client/streamText.ts b/src/vercel/client/streamText.ts index 3b5b11db..443f3482 100644 --- a/src/vercel/client/streamText.ts +++ b/src/vercel/client/streamText.ts @@ -135,6 +135,7 @@ export async function streamText< // When false (saveStreamDeltas.returnImmediately === true), we cannot // defer the final-step save to a post-await block — the function has // already returned by the time onStepFinish fires. See issue #265. + const savesMessages = options?.storageOptions?.saveMessages !== "none"; const willAwaitStream = Boolean(threadId) && (options.saveStreamDeltas === true || @@ -156,10 +157,10 @@ export async function streamText< materialize: (parts) => materializeUIMessageChunkFiles(ctx, component, parts), abortSignal: args.abortSignal, - // Both paths below finish the stream row atomically with the - // message save (issue #181), so the streamer must never finish it - // itself at end-of-stream. - finishHandledExternally: true, + // The message save finishes the stream row atomically (issue + // #181) — but only when there is a save. With saveMessages set to + // "none" nothing does, so the streamer keeps finish ownership. + finishHandledExternally: savesMessages, }, { threadId, @@ -243,6 +244,9 @@ export async function streamText< false, finishStreamId, ); + if (!savesMessages) { + await streamer.finish(); + } initialResponseMessagesSaved = true; } } else { diff --git a/src/vercel/client/streaming.ts b/src/vercel/client/streaming.ts index 27334beb..010d188c 100644 --- a/src/vercel/client/streaming.ts +++ b/src/vercel/client/streaming.ts @@ -218,18 +218,15 @@ export class DeltaStreamer { #cursor: number = 0; public abortController: AbortController; /** - * Who owns the terminal `streams.finish` transition. When true, external code - * finishes the stream row (e.g. atomically with the message save in - * `addMessages`, see issue #181), so `consumeStream` must not finish it. - * Decided once, at construction, by whoever owns the generation. + * When true, external code finishes the stream row (atomically with the + * message save, issue #181) and `consumeStream` must not. Decided once, at + * construction. */ #finishHandledExternally: boolean; /** - * Whether we've stopped accepting new parts. Only meaningful when the stream - * row is finished before the source stream has been fully consumed (the - * `returnImmediately` path). This is an optimization, not a race guard: - * `streams.addDelta` already returns false for a non-streaming row and - * `#sendDelta` treats that as a benign late miss. + * Set only where the row is finished before the source is drained. An + * optimization, not a race guard: `addDelta` already returns false for a + * non-streaming row. */ #stoppedAccepting: boolean = false; @@ -241,12 +238,7 @@ export class DeltaStreamer { onAsyncAbort: (reason: string) => Promise; abortSignal: AbortSignal | undefined; compress: ((parts: T[]) => T[]) | null; - /** - * Set when external code will call `streams.finish` (typically - * atomically with the message save). `consumeStream` then flushes at EOF - * but leaves the terminal transition to that caller. Defaults to false: - * `consumeStream` finishes the stream itself. - */ + /** Defaults to false, meaning `consumeStream` finishes the stream. */ finishHandledExternally?: boolean; materialize?: (parts: T[]) => Promise<{ parts: T[]; @@ -313,16 +305,11 @@ export class DeltaStreamer { if (this.abortController.signal.aborted) { return; } - // Once we've stopped accepting (the returnImmediately path finished the - // stream row inline), the stream record is already "finished" in the DB. - // Late deltas would be silently dropped by streams.addDelta — skip the work. if (this.#stoppedAccepting) { return; } // Buffer before awaiting: a part parked in stream creation would be - // invisible to #flushPendingParts, which only inspects the buffer and the - // in-flight write. The check above and this push are synchronous, so a - // concurrent flush either sees this part or has yet to set the flag. + // invisible to #flushPendingParts. this.#nextParts.push(...parts); await this.getStreamId(); if (this.#stoppedAccepting || this.abortController.signal.aborted) { @@ -356,10 +343,9 @@ export class DeltaStreamer { await this.#waitForAbortCleanup(); return; } - // EOF is the only point at which every part the source will ever hand over - // has actually been handed over — parts live inside the AI SDK's transform - // and tee pipeline until the iterator yields them, so no earlier callback - // can observe them. Drain here, before the terminal transition. + // EOF is the only point where every part has actually been handed over: + // parts sit in the AI SDK's transform and tee pipeline until the iterator + // yields them, so no earlier callback can observe them. try { await this.#flushPendingParts(); } catch (error) { @@ -372,8 +358,6 @@ export class DeltaStreamer { await this.#waitForAbortCleanup(); return; } - // Skip finish if it will be handled externally (atomically with the - // message save). if (!this.#finishHandledExternally) { await this.finish(); } @@ -400,17 +384,13 @@ export class DeltaStreamer { } /** - * Flush buffered parts and stop accepting new ones, for callers that finish - * the stream row before the source stream has been fully consumed — i.e. the - * `returnImmediately` path, where nothing awaits consumption and the save has - * to happen inline (issue #265). + * For the `returnImmediately` path, where nothing awaits consumption and the + * save has to happen inline (issue #265). * - * There is an inherent window here: parts still inside the AI SDK's - * transform/tee pipeline at this instant have not reached addParts, so they - * are never persisted as deltas (the stream-level `finish` chunk among them). - * That is acceptable because the message saved alongside this transition is - * complete and atomic — the saved message, not the delta log, is - * authoritative once the stream is finished. + * Inherent window: parts still inside the AI SDK pipeline at this instant + * never reach addParts, so they are never persisted as deltas (the + * stream-level `finish` chunk among them). The message saved alongside this + * transition is complete, and is authoritative once the stream is finished. */ public async flushAndStopAccepting(): Promise { await this.#flushPendingParts();