fix(consult): preserve streaming across disconnects

frontend/src/lib/stream-text-response.ts

@@ -1,7 +1,7 @@
 type StreamHooks = {
   readonly onFirstOutput?: () => Promise<void>;
-  readonly onComplete?: () => Promise<void>;
-  readonly onError?: (error: unknown, emitted: boolean) => Promise<void>;
+  readonly onComplete?: (output: string) => Promise<void>;
+  readonly onError?: (error: unknown, emitted: boolean, output: string) => Promise<void>;
   readonly onCancel?: (emitted: boolean) => Promise<void>;
 };
 
@@ -10,6 +10,7 @@ type StreamTextResponseOptions = StreamHooks & {
   readonly requestId: string;
   readonly headers?: Record<string, string>;
   readonly transformText?: (text: string) => string;
+  readonly continueAfterDisconnect?: boolean;
 };
 
 const hiddenBlockOpeners = [
@@ -52,8 +53,9 @@ function createVisibleTextTransformer(transform: (text: string) => string) {
             rawBuffer = "";
             hidden = false;
           } else {
-            hiddenBuffer += rawBuffer;
-            rawBuffer = "";
+            const retainedLength = Math.min(2, rawBuffer.length);
+            hiddenBuffer += rawBuffer.slice(0, rawBuffer.length - retainedLength);
+            rawBuffer = rawBuffer.slice(rawBuffer.length - retainedLength);
           }
           break;
         }
@@ -149,16 +151,14 @@ export function streamTextResponse(
 ) {
   const iterator = stream[Symbol.asyncIterator]();
   const encoder = new TextEncoder();
-  // Keep a full natural-language clause unflushed so a later stream chunk cannot
-  // turn an allowed prefix into a disallowed timing or guaranteed conclusion.
-  const guardTailLength = options.transformText ? 1024 : 0;
   const visibleTransformer = options.transformText
     ? createVisibleTextTransformer(options.transformText)
     : null;
-  let pending = "";
   let settled = false;
   let cancellationStarted = false;
+  let disconnected = false;
   let emitted = false;
+  let fullOutput = "";
   let firstOutputSettlementStarted = false;
 
   function startFirstOutputSettlement(value: string) {
@@ -167,67 +167,70 @@ export function streamTextResponse(
     return options.onFirstOutput?.();
   }
 
-  async function enqueueOutput(
-    controller: ReadableStreamDefaultController<Uint8Array>,
+  async function output(
+    controller: ReadableStreamDefaultController<Uint8Array> | undefined,
     value: string,
   ) {
+    if (!value) return;
+    if (disconnected || !controller) {
+      fullOutput += value;
+      return;
+    }
     const firstOutputSettlement = startFirstOutputSettlement(value);
     controller.enqueue(encoder.encode(value));
+    fullOutput += value;
     if (/\S/.test(value)) emitted = true;
     await firstOutputSettlement;
   }
 
-  const body = new ReadableStream<Uint8Array>({
-    async pull(controller) {
-      try {
-        while (true) {
-          const next = await iterator.next();
+  async function consume(
+    controller: ReadableStreamDefaultController<Uint8Array> | undefined,
+  ) {
+    try {
+      while (true) {
+        const { done, value } = await iterator.next();
+        if (settled) return;
+        if (done) {
+          await output(controller, visibleTransformer ? visibleTransformer.finish("") : "");
           if (settled) return;
-          const { done, value } = next;
-          if (done) {
-            const finalText = visibleTransformer
-              ? visibleTransformer.finish(pending)
-              : pending;
-            if (finalText) {
-              await enqueueOutput(controller, finalText);
-            }
-            if (settled) return;
-            settled = true;
-            if (!emitted) {
-              const error = new Error("empty_stream");
-              await options.onError?.(error, false);
-              controller.error(error);
-              return;
-            }
-            await options.onComplete?.();
-            controller.close();
-            return;
-          }
-          pending += value;
-          if (pending.length <= guardTailLength) continue;
-
-          const stableLength = pending.length - guardTailLength;
-          const stable = pending.slice(0, stableLength);
-          pending = pending.slice(stableLength);
-          const transformed = visibleTransformer
-            ? visibleTransformer.push(stable)
-            : stable;
-          if (transformed) {
-            await enqueueOutput(controller, transformed);
-            return;
-          }
-        }
-      } catch (error) {
-        if (cancellationStarted) return;
-        if (!settled) {
           settled = true;
-          await options.onError?.(error, emitted);
+          if (!/\S/.test(fullOutput)) {
+            const error = new Error("empty_stream");
+            await options.onError?.(error, false, fullOutput);
+            if (!disconnected) controller?.error(error);
+            return;
+          }
+          await options.onComplete?.(fullOutput);
+          if (!disconnected) controller?.close();
+          return;
         }
-        controller.error(error);
+
+        const transformed = visibleTransformer
+          ? visibleTransformer.push(value)
+          : value;
+        await output(controller, transformed);
+        if (settled) return;
       }
+    } catch (error) {
+      if (cancellationStarted) return;
+      if (!settled) {
+        settled = true;
+        await options.onError?.(error, emitted, fullOutput);
+      }
+      if (!disconnected) controller?.error(error);
+    }
+  }
+
+  const body = new ReadableStream<Uint8Array>({
+    start(controller) {
+      void consume(controller).catch(() => {});
     },
     async cancel() {
       if (settled) return;
+      if (options.continueAfterDisconnect) {
+        disconnected = true;
+        return;
+      }
       settled = true;
       cancellationStarted = true;
       try {