fix(consult): preserve streaming across disconnects

2026-08-08 21:34:03 +08:00
parent c16c9bfbde
commit 32b683ebef
9 changed files with 1529 additions and 179 deletions
@@ -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) => {