refactor: handle backpressure on stdout
This commit is contained in:
parent
7214fad3e7
commit
bbdb656214
1 changed files with 37 additions and 20 deletions
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue