diff --git a/src/vercel/client/streamText.test.ts b/src/vercel/client/streamText.test.ts index 4a564218..282599d6 100644 --- a/src/vercel/client/streamText.test.ts +++ b/src/vercel/client/streamText.test.ts @@ -113,6 +113,46 @@ 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 }; + }, +}); + export const streamTextCleanupFailure = action({ args: { threadId: v.string() }, handler: async (ctx, { threadId }) => { @@ -153,6 +193,8 @@ export const streamTextCleanupFailure = action({ const testApi: ApiFromModules<{ fns: { streamTextReturnImmediately: typeof streamTextReturnImmediately; + streamTextThrottled: typeof streamTextThrottled; + streamTextThrottledAwaited: typeof streamTextThrottledAwaited; streamTextEmptyAwaited: typeof streamTextEmptyAwaited; streamTextEmptyReturnImmediately: typeof streamTextEmptyReturnImmediately; streamTextCleanupFailure: typeof streamTextCleanupFailure; @@ -317,3 +359,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 f21a13ee..4eea6c3a 100644 --- a/src/vercel/client/streamText.ts +++ b/src/vercel/client/streamText.ts @@ -155,6 +155,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, @@ -225,10 +229,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), @@ -236,7 +243,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({ ifAborted: "returnUndefined", }); @@ -300,6 +310,11 @@ export async function streamText< await call.save(pendingFinalStep, false, finishStreamId); } pendingFinalStep = undefined; + } 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 6a60b91f..7db2a26c 100644 --- a/src/vercel/client/streaming.test.ts +++ b/src/vercel/client/streaming.test.ts @@ -492,4 +492,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 4f42f2c7..25802bdb 100644 --- a/src/vercel/client/streaming.ts +++ b/src/vercel/client/streaming.ts @@ -218,9 +218,18 @@ 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; + /** + * 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; + /** + * 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; constructor( public readonly component: AgentComponent, @@ -230,6 +239,8 @@ export class DeltaStreamer { onAsyncAbort: (reason: string) => Promise; abortSignal: AbortSignal | undefined; compress: ((parts: T[]) => T[]) | null; + /** Defaults to false, meaning `consumeStream` finishes the stream. */ + finishHandledExternally?: boolean; materialize?: (parts: T[]) => Promise<{ parts: T[]; fileRefs: Array<{ url: string; fileId: string }>; @@ -253,6 +264,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) { @@ -294,18 +306,19 @@ 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) { + if (this.#stoppedAccepting) { return; } + // Buffer before awaiting: a part parked in stream creation would be + // invisible to #flushPendingParts. + this.#nextParts.push(...parts); const streamId = await this.getOrCreateStreamId({ ifAborted: "returnUndefined", }); if (!streamId) return; - this.#nextParts.push(...parts); + if (this.#stoppedAccepting || this.abortController.signal.aborted) { + return; + } if ( !this.#ongoingWrite && Date.now() - this.#latestWrite >= this.config.throttleMs @@ -326,24 +339,64 @@ export class DeltaStreamer { await this.#abort(errorToString(error)).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 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) { + if (this.abortController.signal.aborted) { + await this.#waitForAbortCleanup(); + } + throw error; + } + if (this.abortController.signal.aborted) { + await this.#waitForAbortCleanup(); + return; + } + 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(); + } + } + + /** + * For the `returnImmediately` path, where nothing awaits consumption and the + * save has to happen inline (issue #265). + * + * 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 markFinishedExternally(): void { - this.#finishedExternally = true; + public async flushAndStopAccepting(): Promise { + await this.#flushPendingParts(); + this.#stoppedAccepting = true; } /** @@ -378,6 +431,7 @@ export class DeltaStreamer { } let success: boolean; try { + await this.getStreamId(); const delta = await this.#createDelta(); if (!delta) { return; @@ -392,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");