refactor: handle backpressure on stdout

This commit is contained in:
mitchell 2026-04-15 22:36:58 -04:00
parent 7214fad3e7
commit bbdb656214

View file

@ -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);
await checkContainment(part.text, numColumns, numRows);
write(part.text);
checkContainment(part.text);
} else if (part.type === "reasoning-end") {
process.stdout.write("\n\u001b[22m");
await moveCursor(0, -lineCount - 2);
await clearScreenDown();
write("\n\u001b[22m\u001b[23m");
moveCursor(0, -lineCount - 2);
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);
await checkContainment(part.text, numColumns, numRows);
write(part.text);
checkContainment(part.text);
} else if (part.type === "text-end") {
await moveCursor(0, -lineCount - 1);
await clearScreenDown();
moveCursor(0, -lineCount - 1);
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);
await clearScreenDown();
moveCursor(0, -lineCount);
clearScreenDown();
lineCount = 0;
charsSinceNL = 0;
}