refactor: handle backpressure on stdout

src/render/index.ts

@@ -1,10 +1,5 @@
 import type { AsyncIterableStream, TextStreamPart, ToolSet } from "ai";
 
-const moveCursor = (x: number, y: number) =>
-  new Promise<void>((resolve) => process.stdout.moveCursor(x, y, resolve));
-const clearScreenDown = () =>
-  new Promise<void>((resolve) => process.stdout.clearScreenDown(resolve));
-
 let lineCount = 0;
 let charsSinceNL = 0;
 
@@ -13,37 +8,59 @@ function resetCounts() {
   charsSinceNL = 0;
 }
 
+let memOk = true;
+
+process.stdout.addListener("drain", () => {
+  memOk = true;
+});
+
+function write(text: string) {
+  memOk = process.stdout.write(text);
+}
+
+function moveCursor(x: number, y: number) {
+  memOk = process.stdout.moveCursor(x, y);
+}
+
+function clearScreenDown() {
+  memOk = process.stdout.clearScreenDown();
+}
+
 export async function renderStream(
   stream: AsyncIterableStream<TextStreamPart<ToolSet>>,
 ) {
-  const [numColumns, numRows] = process.stdout.getWindowSize();
   resetCounts();
 
   for await (const part of stream) {
     if (part.type === "reasoning-start") {
-      process.stdout.write("\u001b[2m\n");
+      write("\u001b[2m\u001b[3m\n");
     } else if (part.type === "reasoning-delta") {
-      process.stdout.write(part.text);
+      write(part.text);
-      await checkContainment(part.text, numColumns, numRows);
+      checkContainment(part.text);
     } else if (part.type === "reasoning-end") {
-      process.stdout.write("\n\u001b[22m");
+      write("\n\u001b[22m\u001b[23m");
-      await moveCursor(0, -lineCount - 2);
+      moveCursor(0, -lineCount - 2);
-      await clearScreenDown();
+      clearScreenDown();
       resetCounts();
     } else if (part.type === "text-start") {
-      process.stdout.write("\n");
+      write("\n");
     } else if (part.type === "text-delta") {
-      process.stdout.write(part.text);
+      write(part.text);
-      await checkContainment(part.text, numColumns, numRows);
+      checkContainment(part.text);
     } else if (part.type === "text-end") {
-      await moveCursor(0, -lineCount - 1);
+      moveCursor(0, -lineCount - 1);
-      await clearScreenDown();
+      clearScreenDown();
       resetCounts();
     }
+
+    while (!memOk) {
+      await new Promise((resolve) => setTimeout(resolve, 10));
+    }
   }
 }
 
-async function checkContainment(text: string, columns: number, rows: number) {
+function checkContainment(text: string) {
+  const [columns, rows] = process.stdout.getWindowSize();
   charsSinceNL += text.length;
   if (text.includes("\n")) {
     lineCount += text.match(/\n/g)?.length ?? 1;
@@ -54,8 +71,8 @@ async function checkContainment(text: string, columns: number, rows: number) {
   }
 
   if (lineCount >= rows / 2) {
-    await moveCursor(0, -lineCount);
+    moveCursor(0, -lineCount);
-    await clearScreenDown();
+    clearScreenDown();
     lineCount = 0;
     charsSinceNL = 0;
   }