From 2eb978064c27d24f7c0557800d400b03076bc59f Mon Sep 17 00:00:00 2001 From: Jesse_Chen Date: Tue, 21 Jul 2026 13:11:52 +0800 Subject: [PATCH] fix: align stream cancellation with committed output --- frontend/src/lib/stream-text-response.ts | 14 +-- frontend/tests/stream-text-response.test.ts | 107 +++++++++++++++++++- 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/frontend/src/lib/stream-text-response.ts b/frontend/src/lib/stream-text-response.ts index 0566e4cf..203eda41 100644 --- a/frontend/src/lib/stream-text-response.ts +++ b/frontend/src/lib/stream-text-response.ts @@ -158,21 +158,22 @@ export function streamTextResponse( let pending = ""; let settled = false; let emitted = false; - let firstOutputSettled = false; + let firstOutputSettlementStarted = false; - async function settleFirstOutput(value: string) { - if (!/\S/.test(value) || firstOutputSettled) return; - firstOutputSettled = true; - await options.onFirstOutput?.(); + function startFirstOutputSettlement(value: string) { + if (!/\S/.test(value) || firstOutputSettlementStarted) return undefined; + firstOutputSettlementStarted = true; + return options.onFirstOutput?.(); } async function enqueueOutput( controller: ReadableStreamDefaultController, value: string, ) { - await settleFirstOutput(value); + const firstOutputSettlement = startFirstOutputSettlement(value); controller.enqueue(encoder.encode(value)); if (/\S/.test(value)) emitted = true; + await firstOutputSettlement; } const body = new ReadableStream({ @@ -187,6 +188,7 @@ export function streamTextResponse( if (finalText) { await enqueueOutput(controller, finalText); } + if (settled) return; settled = true; if (!emitted) { const error = new Error("empty_stream"); diff --git a/frontend/tests/stream-text-response.test.ts b/frontend/tests/stream-text-response.test.ts index c7723b32..eaa97d7c 100644 --- a/frontend/tests/stream-text-response.test.ts +++ b/frontend/tests/stream-text-response.test.ts @@ -2,7 +2,7 @@ import assert from "node:assert/strict"; import test from "node:test"; import { streamTextResponse } from "../src/lib/stream-text-response.ts"; -test("durable settlement runs before the first response bytes are exposed", async () => { +test("durable settlement starts before the first response bytes are exposed", async () => { const order: string[] = []; async function* reply() { yield "第一段"; @@ -26,6 +26,111 @@ test("durable settlement runs before the first response bytes are exposed", asyn assert.deepEqual(order, ["settled", "第一段", "completed"]); }); +test("cancelling while first-output settlement is pending observes committed output", async () => { + let markSettlementStarted = () => {}; + const settlementStarted = new Promise((resolve) => { + markSettlementStarted = resolve; + }); + let releaseSettlement = () => {}; + const settlementGate = new Promise((resolve) => { + releaseSettlement = resolve; + }); + let markSettlementFinished = () => {}; + const settlementFinished = new Promise((resolve) => { + markSettlementFinished = resolve; + }); + async function* reply() { + yield "已经通过转换的可见短回答。"; + } + let cancelledAsEmitted: boolean | null = null; + let completed = 0; + const response = streamTextResponse(reply(), { + mode: "mastra", + requestId: "00000000-0000-4000-8000-000000000100", + transformText: (text) => text, + onFirstOutput: async () => { + markSettlementStarted(); + try { + await settlementGate; + } finally { + markSettlementFinished(); + } + }, + onCancel: async (emitted) => { cancelledAsEmitted = emitted; }, + onComplete: async () => { completed += 1; }, + }); + const reader = response.body?.getReader(); + assert.ok(reader); + const pendingRead = reader.read(); + await settlementStarted; + + try { + await reader.cancel(); + const first = await pendingRead; + assert.equal(cancelledAsEmitted, true); + assert.equal(first.done, false); + assert.match(new TextDecoder().decode(first.value), /已经通过转换的可见短回答/); + } finally { + releaseSettlement(); + } + await settlementFinished; + await new Promise((resolve) => queueMicrotask(resolve)); + assert.equal(completed, 0); +}); + +test("a synchronous first-output hook failure exposes no bytes", async () => { + let observedEmitted: boolean | null = null; + async function* reply() { + yield "已经通过转换的可见正文。".repeat(120); + } + const response = streamTextResponse(reply(), { + mode: "mastra", + requestId: "00000000-0000-4000-8000-000000000101", + transformText: (text) => text, + onFirstOutput: () => { throw new Error("settlement_start_failed"); }, + onError: async (_error, emitted) => { observedEmitted = emitted; }, + }); + + await assert.rejects(response.text(), /settlement_start_failed/); + assert.equal(observedEmitted, false); +}); + +test("an asynchronous first-output rejection preserves already committed bytes", async () => { + let rejectSettlement = () => {}; + const settlement = new Promise((_resolve, reject) => { + rejectSettlement = () => reject(new Error("settlement_failed")); + }); + let markSettlementStarted = () => {}; + const settlementStarted = new Promise((resolve) => { + markSettlementStarted = resolve; + }); + async function* reply() { + yield "已经通过转换的可见正文。".repeat(120); + } + let observedEmitted: boolean | null = null; + const response = streamTextResponse(reply(), { + mode: "mastra", + requestId: "00000000-0000-4000-8000-000000000102", + transformText: (text) => text, + onFirstOutput: () => { + markSettlementStarted(); + return settlement; + }, + onError: async (_error, emitted) => { observedEmitted = emitted; }, + }); + const reader = response.body?.getReader(); + assert.ok(reader); + const pendingRead = reader.read(); + await settlementStarted; + rejectSettlement(); + + const first = await pendingRead; + assert.equal(first.done, false); + assert.match(new TextDecoder().decode(first.value), /已经通过转换的可见正文/); + await assert.rejects(reader.read(), /settlement_failed/); + assert.equal(observedEmitted, true); +}); + test("charges a consultation when cancellation happens after partial output", async () => { // Given let completed = 0;