fix(consult): preserve streaming across disconnects

frontend/tests/application-billing-contract.test.ts

@@ -30,13 +30,13 @@ test("standard consultation resolves and settles the session-pinned model versio
   assert.match(consultRoute, /modelConfigVersion: selectedModel\.configVersion/);
 });
 
-test("standard consultation awaits real usage before its only permanent settlement", () => {
+test("standard consultation awaits real usage before durable response settlement", () => {
   assert.doesNotMatch(consultRoute, /recordActualUsage|void usage\.then/);
   assert.doesNotMatch(consultRoute, /inputTokens: 0,[\s\S]*outputTokens: 0,[\s\S]*costMicrousd: 0/);
-  assert.match(consultRoute, /async function complete\(usage: Promise<\{ inputTokens\?: number; outputTokens\?: number \}>\)/);
+  assert.match(consultRoute, /async function usagePayload\(usage: Promise<\{ inputTokens\?: number; outputTokens\?: number \}>\)/);
   assert.match(consultRoute, /const resolved = await usage;/);
-  assert.match(consultRoute, /await completeUsage\(accounting, userId, requestId, \{[\s\S]*inputTokens,[\s\S]*outputTokens,[\s\S]*costMicrousd/);
-  assert.match(consultRoute, /const completeWithUsage = \(\) => complete\(result\.totalUsage\)/g);
+  assert.match(consultRoute, /const actualUsage = await usagePayload\(usage\);[\s\S]*p_actual_usage: actualUsage/);
+  assert.equal(consultRoute.match(/result\.totalUsage/g)?.length, 4);
 });
 
 test("standard consultation forwards its stable reservation request as the usage event key", async () => {
@@ -57,7 +57,8 @@ test("standard consultation forwards its stable reservation request as the usage
   await completeUsage(accounting, "00000000-0000-4000-8000-000000000001", eventKey, usage);
 
   assert.deepEqual(calls.map((call) => (call.p_actual_usage as { eventKey: string }).eventKey), [eventKey, eventKey]);
-  assert.match(consultRoute, /completeUsage\(accounting, userId, requestId, \{[\s\S]*eventKey: requestId,/);
+  assert.match(consultRoute, /eventKey: requestId,/);
+  assert.match(consultRoute, /const actualUsage = await usagePayload\(usage\);[\s\S]*p_actual_usage: actualUsage/);
   assert.doesNotMatch(consultRoute, /eventKey:\s*(?:globalThis\.)?crypto\.randomUUID\(\)/);
 });
 

frontend/tests/consultation-recovery.test.ts

@@ -0,0 +1,120 @@
+import assert from "node:assert/strict";
+import { readFileSync } from "node:fs";
+import test from "node:test";
+
+const source = readFileSync(new URL("../src/app/page.tsx", import.meta.url), "utf8");
+const sendSource = source.slice(source.indexOf("  async function send("), source.indexOf("\n\n  useGSAP", source.indexOf("  async function send(")));
+const stopSource = source.slice(source.indexOf("  async function stopResponse("), source.indexOf("\n\n  function completeConsultationInterface"));
+
+test("consultation persists the optimistic user message before generation starts", () => {
+  const undo = sendSource.indexOf("await waitForUndoWindow(controller.signal)");
+  const persist = sendSource.indexOf("await persistSession(userSession)");
+  const consult = sendSource.indexOf('fetch("/api/consult"');
+
+  assert.ok(undo >= 0 && undo < persist && persist < consult);
+  assert.match(sendSource, /问题保存失败，未开始生成；问题已放回输入框。[\s\S]*?return false;[\s\S]*?fetch\("\/api\/consult"/);
+});
+
+test("only explicit stop can request consultation cancellation", () => {
+  const streamCatch = sendSource.indexOf("    } catch (caught) {", sendSource.indexOf('fetch("/api/consult"'));
+
+  assert.equal(source.match(/confirmCancellation\(/g)?.length, 2);
+  assert.match(stopSource, /await confirmCancellation\(/);
+  assert.ok(streamCatch >= 0);
+  assert.doesNotMatch(sendSource.slice(streamCatch), /confirmCancellation\(|\/api\/consult\/cancel/);
+});
+
+test("durable partial stop cancels before preserving content and recovers completed conflicts", () => {
+  const partialStop = stopSource.slice(
+    stopSource.indexOf("    if (pending.partialReply)"),
+    stopSource.indexOf("\n    updateSession(pending.sessionId, () => pending.previousSession)"),
+  );
+  const cancel = partialStop.indexOf("await requestCancellation(pending.requestId)");
+  const persist = partialStop.indexOf("await persistSession(stoppedSession)");
+  const conflict = partialStop.indexOf("error.status === 409");
+  const recoveryPersist = partialStop.slice(conflict, persist);
+
+  assert.ok(cancel >= 0 && cancel < persist);
+  assert.ok(conflict >= 0 && conflict < persist);
+  assert.doesNotMatch(recoveryPersist, /persistSession\(stoppedSession\)/);
+  assert.match(recoveryPersist, /fetchConsultationStatus\(pending.sessionId, pending.requestId\)/);
+  assert.match(recoveryPersist, /status.status === "completed"[\s\S]*?fetchSessions\(\)[\s\S]*?readSessions\(payload, modelCatalog\)/);
+  assert.match(partialStop, /已停止回答，现有内容已保留，本次点数已退回。/);
+  assert.match(source, /停止回答，保留已生成内容并退回本次点数/);
+  assert.match(source, /停止回答，保留现有内容并申请退回本次点数/);
+});
+
+test("explicit consultation HTTP failures unlock instead of entering recovery", () => {
+  assert.match(sendSource, /!response\.ok[\s\S]*throw new ConsultationResponseError\([\s\S]*response\.status/);
+  const explicitFailure = sendSource.slice(
+    sendSource.indexOf("caught instanceof ConsultationResponseError"),
+    sendSource.indexOf("if (!cancelled && ownsInterface && pendingConsultation.current)"),
+  );
+  assert.match(explicitFailure, /setRequestError\(\{ sessionId, message: caught\.message \}\)/);
+  assert.match(explicitFailure, /completeConsultationInterface\(requestId\)/);
+  assert.doesNotMatch(explicitFailure, /phase: "recovering"|setConsultationPhase\("recovering"\)/);
+});
+
+test("reserved consultations recover through the status endpoint", () => {
+  assert.match(source, /\/api\/consult\/status\?sessionId=\$\{encodeURIComponent\(sessionId\)\}&requestId=\$\{encodeURIComponent\(requestId\)\}/);
+  assert.match(source, /fetch\("\/api\/consult\/status", \{ signal, cache: "no-store" \}\)/);
+  assert.match(source, /reservedConsultation\?\.status === "reserved"[\s\S]*nextSessions\.find\(\(session\) => session\.id === reservedConsultation\.sessionId\)/);
+  assert.match(source, /status\?\.status !== "reserved"[\s\S]*sessions\.find\(\(item\) => item\.id === status\.sessionId\)/);
+  assert.match(source, /status: "reserved" \| "completed" \| "cancelled"/);
+  assert.match(source, /readonly responseMessage\?: unknown/);
+  assert.match(source, /phase: "recovering"/);
+  assert.match(source, /window\.setTimeout\(\(\) => void poll\(\), 1_750\)/);
+  assert.match(source, /status\.status === "completed"[\s\S]*?fetchSessions\(controller\.signal\)[\s\S]*?readSessions\(payload, modelCatalog\)/);
+  assert.match(source, /window\.addEventListener\("online", onOnline\)/);
+  assert.match(source, /window\.addEventListener\("pageshow", onPageShow\)/);
+  assert.match(source, /网络已断开，回答仍在后台生成；联网后会自动恢复。/);
+});
+
+test("three consecutive strict status misses unlock recovery while transient failures keep polling", () => {
+  assert.match(source, /class ConsultationStatusError extends Error[\s\S]*readonly status: number/);
+  assert.match(source, /throw new ConsultationStatusError\([\s\S]*response\.status/);
+  const recoveryEffect = source.slice(
+    source.indexOf('if (consultationPhase !== "recovering"'),
+    source.indexOf("}, [consultationPhase, modelCatalog, pendingRequestId, pendingSessionId])"),
+  );
+  assert.match(source, /const consultationStatusMissingCount = useRef\(0\)/);
+  assert.match(recoveryEffect, /status\.status === "reserved"[\s\S]*consultationStatusMissingCount\.current = 0/);
+  assert.match(recoveryEffect, /caught instanceof ConsultationStatusError && caught\.status === 404[\s\S]*consultationStatusMissingCount\.current \+= 1[\s\S]*consultationStatusMissingCount\.current >= 3/);
+  assert.match(recoveryEffect, /setPendingSessionId\(null\)[\s\S]*setPendingRequestId\(null\)[\s\S]*setConsultationPhase\(null\)/);
+  assert.match(recoveryEffect, /后台未找到本次咨询请求，已停止恢复，请重新发送。/);
+  assert.match(recoveryEffect, /consultationStatusMissingCount\.current = 0;[\s\S]*回答仍在后台生成，正在自动恢复。/);
+  assert.match(recoveryEffect, /consultationStatusMissingCount\.current > 0[\s\S]*window\.setTimeout\(\(\) => void poll\(\), 1_750\)[\s\S]*else \{[\s\S]*void poll\(\)/);
+});
+
+test("tab-local pending ids drive strict bootstrap recovery before the global fallback", () => {
+  const bootstrap = source.slice(
+    source.indexOf("let reservedConsultation: ConsultationStatus | null = null"),
+    source.indexOf("if (controller.signal.aborted) return;", source.indexOf("let reservedConsultation: ConsultationStatus | null = null")),
+  );
+  const storageSync = source.slice(
+    source.indexOf("if (!hydrated || uiPreview.current) return;", source.indexOf("}, []);")),
+    source.indexOf('if (consultationPhase !== "recovering"'),
+  );
+
+  assert.match(bootstrap, /sessionStorage\.getItem\(pendingConsultationStorageKey\)/);
+  assert.match(bootstrap, /uuidPattern\.test\(parsedPending\.sessionId\)[\s\S]*uuidPattern\.test\(parsedPending\.requestId\)[\s\S]*nextSessions\.some\(\(session\) => session\.id === parsedPending\.sessionId\)/);
+  assert.match(bootstrap, /else \{[\s\S]*sessionStorage\.removeItem\(pendingConsultationStorageKey\)[\s\S]*\} catch \{[\s\S]*sessionStorage\.removeItem\(pendingConsultationStorageKey\)/);
+  assert.ok(bootstrap.indexOf("if (storedPending)") < bootstrap.indexOf("fetchActiveConsultationStatus(controller.signal)"));
+  assert.match(bootstrap, /if \(storedPending\) \{[\s\S]*fetchConsultationStatus\([\s\S]*storedPending\.sessionId,[\s\S]*storedPending\.requestId,[\s\S]*\} else \{[\s\S]*fetchActiveConsultationStatus/);
+  assert.match(bootstrap, /status\.status === "reserved"[\s\S]*reservedConsultation = status;[\s\S]*else \{[\s\S]*sessionStorage\.removeItem\(pendingConsultationStorageKey\)/);
+  assert.match(bootstrap, /caught instanceof ConsultationStatusError && caught\.status === 404 \? 1 : 0/);
+  assert.match(storageSync, /pendingSessionId && pendingRequestId[\s\S]*sessionStorage\.setItem\(pendingConsultationStorageKey,[\s\S]*sessionId: pendingSessionId,[\s\S]*requestId: pendingRequestId/);
+  assert.match(storageSync, /else \{[\s\S]*sessionStorage\.removeItem\(pendingConsultationStorageKey\)/);
+});
+
+test("the first default consultation title is persisted with the user question", () => {
+  const userSessionBlock = sendSource.slice(
+    sendSource.indexOf("const userSession: ChatSession"),
+    sendSource.indexOf("const requestId = globalThis.crypto.randomUUID()"),
+  );
+  assert.match(userSessionBlock, /currentSession\.messages\.length === 0 && currentSession\.title === "新对话"[\s\S]*resolveSessionTitle\(question\)/);
+  assert.ok(sendSource.indexOf("await persistSession(userSession)") < sendSource.indexOf('fetch("/api/consult"'));
+  assert.doesNotMatch(sendSource, /persistSession\(completedSession\)/);
+  assert.match(sendSource, /const completedSession: ChatSession = \{[\s\S]*title: userSession\.title/);
+  assert.doesNotMatch(sendSource, /resolveSessionTitle\(question, reply\.title\)/);
+});

frontend/tests/consultation-stream-recovery.test.ts

@@ -0,0 +1,98 @@
+import assert from "node:assert/strict";
+import { readFileSync } from "node:fs";
+import test from "node:test";
+
+const read = (path: string) => readFileSync(new URL(`../${path}`, import.meta.url), "utf8");
+const consultRoute = read("src/app/api/consult/route.ts");
+const statusRoute = read("src/app/api/consult/status/route.ts");
+const migration = read("supabase/migrations/20260808030000_consultation_stream_recovery.sql");
+
+test("reserves usage and binds the owned consultation session atomically", () => {
+  assert.match(migration, /add column if not exists session_id uuid references public\.chat_sessions\(id\) on delete set null/i);
+  assert.match(migration, /reserve_consultation_usage\([\s\S]*p_session_id uuid[\s\S]*from public\.chat_sessions as session[\s\S]*session\.id = p_session_id[\s\S]*session\.user_id = p_user_id[\s\S]*session\.session_type = 'consultation'/i);
+  assert.match(migration, /from public\.authorize_usage[\s\S]*insert into public\.consultation_requests\(user_id, request_id, session_id, status\)[\s\S]*values \(p_user_id, btrim\(p_request_id\), p_session_id, 'reserved'\)/i);
+  assert.match(migration, /reserve_consultation_usage\(uuid, text, uuid, text, integer\)/i);
+  assert.match(consultRoute, /rpc\("reserve_consultation_usage"[\s\S]*p_session_id: sessionId/);
+  assert.doesNotMatch(consultRoute, /from\("consultation_requests"\)[\s\S]*\.update\(\{ session_id: sessionId/);
+});
+
+test("persists transformed assistant metadata before atomically settling usage", () => {
+  assert.equal(consultRoute.match(/continueAfterDisconnect: true/g)?.length, 2);
+  assert.equal(consultRoute.match(/onComplete: \(rawTransformedText\) => settle\(\(\) => completeResponse\(/g)?.length, 2);
+  assert.match(consultRoute, /parseAgentReply\(rawTransformedText, consultationTheme\)/);
+  assert.match(consultRoute, /role: "assistant" as const,[\s\S]*suggestions: reply\.suggestions,[\s\S]*techniqueTruth,[\s\S]*workflowReceipt/);
+
+  const append = migration.indexOf("set messages = session.messages || jsonb_build_array(p_response_message)");
+  const store = migration.indexOf("set response_message = p_response_message");
+  const settle = migration.indexOf("from public.complete_usage");
+  const complete = migration.indexOf("set status = 'completed'");
+  assert.ok(append >= 0 && append < store && store < settle && settle < complete);
+});
+
+test("persists partial transformed output when the upstream stream errors", () => {
+  assert.doesNotMatch(consultRoute, /completeInterrupted|completeUsage\(accounting, userId, requestId/);
+  assert.equal(
+    consultRoute.match(/onError: \(_error, emitted, output: string\) => settleErrored\(emitted, output\)/g)?.length,
+    2,
+  );
+  assert.equal(consultRoute.match(/const settleErrored = \(emitted: boolean, output: string\) => settle\(/g)?.length, 2);
+  assert.equal(consultRoute.match(/emitted[\s\S]*?\? \(\) => completeResponse\([\s\S]*?output,[\s\S]*?result\.totalUsage,[\s\S]*?: cancel,/g)?.length, 2);
+  assert.equal(consultRoute.match(/onCancel: \(\) => settle\(cancel\)/g)?.length, 2);
+});
+
+test("best-effort cancels a failed or uncertain durable completion before rethrowing", () => {
+  const completion = consultRoute.slice(
+    consultRoute.indexOf("async function completeResponse("),
+    consultRoute.indexOf("let settlement:"),
+  );
+  assert.match(completion, /try \{[\s\S]*accounting\.rpc\("complete_consultation_response"/);
+  assert.match(completion, /catch \(error\) \{[\s\S]*await cancel\(\);[\s\S]*throw error;/);
+  assert.ok(completion.indexOf("await cancel();") > completion.indexOf('accounting.rpc("complete_consultation_response"'));
+});
+
+test("serializes explicit cancellation against response persistence and charging", () => {
+  const cancelFunction = migration.slice(
+    migration.indexOf("create or replace function public.cancel_consultation_credit"),
+    migration.indexOf("create or replace function public.complete_consultation_response"),
+  );
+  assert.match(cancelFunction, /pg_advisory_xact_lock\(hashtextextended\(p_user_id::text \|\| ':' \|\| btrim\(p_request_id\), 0\)\)/);
+  assert.match(cancelFunction, /from public\.consultation_requests[\s\S]*for update/);
+  assert.match(cancelFunction, /if v_request\.status = 'completed' then[\s\S]*'request_completed'/i);
+  assert.match(cancelFunction, /if v_request\.status = 'cancelled' then[\s\S]*select true/i);
+  assert.match(cancelFunction, /if v_request\.status <> 'reserved' then[\s\S]*from public\.release_usage[\s\S]*set status = 'cancelled'/i);
+  assert.match(cancelFunction, /grant execute on function public\.cancel_consultation_credit\(uuid, text\)[\s\S]*to service_role/);
+  assert.match(migration, /if v_request\.status = 'cancelled' then[\s\S]*'request_cancelled'/i);
+  assert.match(consultRoute, /cancel_consultation_credit/);
+  assert.match(consultRoute, /error_code !== "request_cancelled"/);
+});
+
+test("status endpoint supports one global reserved lookup and strict bound polling", () => {
+  assert.ok(statusRoute.indexOf("supabase.auth.getUser()") < statusRoute.indexOf("createAdminSupabaseClient()"));
+  assert.match(statusRoute, /activeLookup = sessionIdValue === null && requestIdValue === null/);
+  assert.match(statusRoute, /activeLookup[\s\S]*\.eq\("status", "reserved"\)\.order\("created_at", \{ ascending: false \}\)\.limit\(1\)/);
+  assert.match(statusRoute, /\.eq\("user_id", user\.id\)[\s\S]*\.eq\("session_id", sessionId!\.data\)[\s\S]*\.eq\("request_id", requestId!\.data\)/);
+  assert.match(statusRoute, /requestId: statusData\.request_id,[\s\S]*sessionId: statusData\.session_id,[\s\S]*status: statusData\.status/);
+  assert.doesNotMatch(statusRoute, /String\(data\.response_message\)|responseMessage:\s*data\.response_message as string/);
+});
+
+test("detached completion and cancellation use a bounded retry ceiling", () => {
+  assert.match(consultRoute, /const detachedSettlementAttempts = 3/);
+  assert.match(consultRoute, /ponytail: Staging MVP ceiling—without a queue\/worker/);
+  assert.match(consultRoute, /for \(let attempt = 1; attempt <= detachedSettlementAttempts; attempt \+= 1\)/);
+  assert.match(consultRoute, /setTimeout\(resolve, attempt \* 150\)/);
+  assert.match(consultRoute, /async function cancel\(\)[\s\S]*retryDetachedSettlement\(async \(\) =>[\s\S]*rpc\("cancel_consultation_credit"/);
+  assert.match(consultRoute, /async function completeResponse\([\s\S]*retryDetachedSettlement\(async \(\) =>[\s\S]*rpc\("complete_consultation_response"/);
+});
+
+test("status expires stale reservations after fifteen minutes but leaves fresh reservations active", () => {
+  assert.match(statusRoute, /const reservedLeaseMs = 15 \* 60 \* 1000/);
+  assert.match(statusRoute, /ponytail: Staging MVP ceiling—without a queue\/worker/);
+  assert.match(statusRoute, /Number\.isFinite\(timestamp\) && now - timestamp >= reservedLeaseMs/);
+  const leaseBlock = statusRoute.slice(
+    statusRoute.indexOf('if (statusData.status === "reserved" && reservationLeaseExpired(statusData.updated_at))'),
+    statusRoute.indexOf("return NextResponse.json({\n    requestId: statusData.request_id"),
+  );
+  assert.match(leaseBlock, /runCreditRpc\([\s\S]*"cancel_consultation_credit"/);
+  assert.match(leaseBlock, /\.eq\("session_id", statusData\.session_id\)[\s\S]*\.eq\("request_id", statusData\.request_id\)/);
+  assert.doesNotMatch(statusRoute.slice(0, statusRoute.indexOf('if (statusData.status === "reserved"')), /cancel_consultation_credit/);
+});

frontend/tests/stream-text-response.test.ts

@@ -26,6 +26,47 @@ test("durable settlement starts before the first response bytes are exposed", as
   assert.deepEqual(order, ["settled", "第一段", "completed"]);
 });
 
+test("transformed output emits at the first complete clause without reading ahead", async () => {
+  let markSecondReadStarted = () => {};
+  const secondReadStarted = new Promise<"read-ahead">((resolve) => {
+    markSecondReadStarted = () => resolve("read-ahead");
+  });
+  const never = new Promise<IteratorResult<string>>(() => {});
+  let reads = 0;
+  const reply: AsyncIterable<string> = {
+    [Symbol.asyncIterator]() {
+      return {
+        next() {
+          reads += 1;
+          if (reads === 1) {
+            return Promise.resolve({ done: false, value: "第一句。" });
+          }
+          markSecondReadStarted();
+          return never;
+        },
+        return() {
+          return Promise.resolve({ done: true, value: undefined });
+        },
+      };
+    },
+  };
+  const response = streamTextResponse(reply, {
+    mode: "mastra",
+    requestId: "00000000-0000-4000-8000-000000000104",
+    transformText: (text) => text,
+  });
+  const reader = response.body?.getReader();
+  assert.ok(reader);
+
+  const first = await Promise.race([reader.read(), secondReadStarted]);
+
+  assert.notEqual(first, "read-ahead");
+  if (first === "read-ahead") assert.fail("read-ahead");
+  assert.equal(first.done, false);
+  assert.equal(new TextDecoder().decode(first.value), "第一句。");
+  await reader.cancel();
+});
+
 test("cancelling while first-output settlement is pending observes committed output", async () => {
   let markSettlementStarted = () => {};
   const settlementStarted = new Promise<void>((resolve) => {
@@ -215,6 +256,192 @@ test("charges a consultation when cancellation happens after partial output", as
   assert.equal(completed, 1);
 });
 
+test("opt-in disconnect keeps consuming and completes with the full transformed output", async () => {
+  let releaseSecondChunk = () => {};
+  const secondChunk = new Promise<void>((resolve) => {
+    releaseSecondChunk = resolve;
+  });
+  let reads = 0;
+  let returnCalls = 0;
+  const reply: AsyncIterable<string> = {
+    [Symbol.asyncIterator]() {
+      return {
+        async next() {
+          reads += 1;
+          if (reads === 1) return { done: false, value: "第一段回答。".repeat(200) };
+          if (reads === 2) {
+            await secondChunk;
+            return { done: false, value: "第二段回答。" };
+          }
+          return { done: true, value: undefined };
+        },
+        return() {
+          returnCalls += 1;
+          return Promise.resolve({ done: true, value: undefined });
+        },
+      };
+    },
+  };
+  let cancelCalls = 0;
+  let completedOutput = "";
+  let markCompleted = () => {};
+  const completed = new Promise<"completed">((resolve) => {
+    markCompleted = () => resolve("completed");
+  });
+  let markCancelled = () => {};
+  const cancelled = new Promise<"cancelled">((resolve) => {
+    markCancelled = () => resolve("cancelled");
+  });
+  const response = streamTextResponse(reply, {
+    mode: "mastra",
+    requestId: "00000000-0000-4000-8000-000000000105",
+    transformText: (text) => text.replaceAll("回答", "安全回答"),
+    continueAfterDisconnect: true,
+    onCancel: async () => {
+      cancelCalls += 1;
+      markCancelled();
+    },
+    onComplete: async (output) => {
+      completedOutput = output;
+      markCompleted();
+    },
+  });
+  const reader = response.body?.getReader();
+  assert.ok(reader);
+  const first = await reader.read();
+  assert.equal(first.done, false);
+
+  await reader.cancel();
+  releaseSecondChunk();
+  assert.equal(await Promise.race([completed, cancelled]), "completed");
+
+  assert.equal(returnCalls, 0);
+  assert.equal(cancelCalls, 0);
+  assert.equal(
+    completedOutput,
+    `${"第一段安全回答。".repeat(200)}第二段安全回答。`,
+  );
+});
+
+test("starts draining without a reader and completes with the full transformed output", { timeout: 1_000 }, async () => {
+  let markFirstReadStarted = () => {};
+  const firstReadStarted = new Promise<void>((resolve) => {
+    markFirstReadStarted = resolve;
+  });
+  let releaseSecondChunk = () => {};
+  const secondChunk = new Promise<void>((resolve) => {
+    releaseSecondChunk = resolve;
+  });
+  let reads = 0;
+  const reply: AsyncIterable<string> = {
+    [Symbol.asyncIterator]() {
+      return {
+        async next() {
+          reads += 1;
+          if (reads === 1) {
+            markFirstReadStarted();
+            return { done: false, value: "第一段回答。" };
+          }
+          if (reads === 2) {
+            await secondChunk;
+            return { done: false, value: "第二段回答。" };
+          }
+          return { done: true, value: undefined };
+        },
+      };
+    },
+  };
+  let completedOutput = "";
+  let markCompleted = () => {};
+  const completed = new Promise<void>((resolve) => {
+    markCompleted = resolve;
+  });
+
+  const response = streamTextResponse(reply, {
+    mode: "mastra",
+    requestId: "00000000-0000-4000-8000-000000000107",
+    transformText: (text) => text.replaceAll("回答", "安全回答"),
+    continueAfterDisconnect: true,
+    onComplete: async (output) => {
+      completedOutput = output;
+      markCompleted();
+    },
+  });
+  assert.ok(response.body);
+  await firstReadStarted;
+
+  releaseSecondChunk();
+  await completed;
+
+  assert.equal(reads, 3);
+  assert.equal(completedOutput, "第一段安全回答。第二段安全回答。");
+});
+
+test("an early opt-in cancel keeps the started producer draining to completion", { timeout: 1_000 }, async () => {
+  let markFirstReadStarted = () => {};
+  const firstReadStarted = new Promise<void>((resolve) => {
+    markFirstReadStarted = resolve;
+  });
+  let releaseFirstChunk = () => {};
+  const firstChunk = new Promise<void>((resolve) => {
+    releaseFirstChunk = resolve;
+  });
+  let reads = 0;
+  let returnCalls = 0;
+  const reply: AsyncIterable<string> = {
+    [Symbol.asyncIterator]() {
+      return {
+        async next() {
+          reads += 1;
+          if (reads === 1) {
+            markFirstReadStarted();
+            await firstChunk;
+            return { done: false, value: "第一段回答。" };
+          }
+          if (reads === 2) return { done: false, value: "第二段回答。" };
+          return { done: true, value: undefined };
+        },
+        return() {
+          returnCalls += 1;
+          return Promise.resolve({ done: true, value: undefined });
+        },
+      };
+    },
+  };
+  let cancelCalls = 0;
+  let completeCalls = 0;
+  let completedOutput = "";
+  let markCompleted = () => {};
+  const completed = new Promise<void>((resolve) => {
+    markCompleted = resolve;
+  });
+  const response = streamTextResponse(reply, {
+    mode: "mastra",
+    requestId: "00000000-0000-4000-8000-000000000108",
+    transformText: (text) => text.replaceAll("回答", "安全回答"),
+    continueAfterDisconnect: true,
+    onCancel: async () => { cancelCalls += 1; },
+    onComplete: async (output) => {
+      completeCalls += 1;
+      completedOutput = output;
+      markCompleted();
+    },
+  });
+  const reader = response.body?.getReader();
+  assert.ok(reader);
+  await firstReadStarted;
+
+  await reader.cancel();
+  releaseFirstChunk();
+  await completed;
+
+  assert.equal(reads, 3);
+  assert.equal(returnCalls, 0);
+  assert.equal(cancelCalls, 0);
+  assert.equal(completeCalls, 1);
+  assert.equal(completedOutput, "第一段安全回答。第二段安全回答。");
+});
+
 test("refunds when cancellation happens before any output", async () => {
   // Given
   let completed = 0;
@@ -342,6 +569,27 @@ test("a transformed short reply that fails while buffered reports no emitted out
   assert.equal(observedEmitted, false);
 });
 
+test("an iterator error after partial output reports the full transformed output", async () => {
+  let observedOutput = "";
+  async function* reply() {
+    yield "第一段回答。";
+    yield "第二段回答。";
+    throw new Error("upstream_failed_after_output");
+  }
+  const response = streamTextResponse(reply(), {
+    mode: "mastra",
+    requestId: "00000000-0000-4000-8000-000000000106",
+    transformText: (text) => text.replaceAll("回答", "安全回答"),
+    onError: async (_error, emitted, output) => {
+      assert.equal(emitted, true);
+      observedOutput = output;
+    },
+  });
+
+  await assert.rejects(response.text(), /upstream_failed_after_output/);
+  assert.equal(observedOutput, "第一段安全回答。第二段安全回答。");
+});
+
 test("cancelling while transformed short output is buffered reports no emitted output", async () => {
   let markSecondReadStarted = () => {};
   const secondReadStarted = new Promise<void>((resolve) => {