lmc/src/render/index.ts

import type { WriteStream } from "node:tty";
import type {
  AsyncIterableStream,
  LanguageModelUsage,
  TextStreamPart,
  ToolSet,
} from "ai";
import styles from "ansi-styles";
import { chalkStderr, supportsColorStderr } from "chalk";
import { renderMarkdown } from "./markdown";

export default class Renderer {
  private lineCount = 0;
  private charsSinceNL = 0;
  private memOk = true;
  private buffer = "";
  private unsafeIndex = 0;

  constructor(
    private isPlain = false,
    private isPlainErr = false,
  ) {
    process.stdout.addListener("drain", () => {
      this.memOk = true;
    });
    process.stderr.addListener("drain", () => {
      this.memOk = true;
    });
  }

  private write(text: string, pipe: WriteStream = process.stdout) {
    this.memOk = pipe.write(text);
  }

  private moveClearReset(y: number, pipe: WriteStream = process.stdout) {
    this.memOk = pipe.moveCursor(-9999, y);
    this.memOk = pipe.clearScreenDown();
    this.resetCounts();
  }

  private resetCounts() {
    this.lineCount = 0;
    this.charsSinceNL = 0;
  }

  private checkContainment(text: string, pipe: WriteStream = process.stdout) {
    const [columns, rows] = pipe.getWindowSize();
    this.charsSinceNL += text.length;
    if (text.includes("\n")) {
      this.lineCount += text.match(/\n/g)?.length ?? 1;
      this.charsSinceNL = text.length - text.lastIndexOf("\n") - 1;
    } else if (this.charsSinceNL > columns) {
      this.lineCount++;
      this.charsSinceNL = this.charsSinceNL - columns;
    }

    if (this.lineCount >= Math.max(rows - 10, 1)) {
      this.moveClearReset(-this.lineCount, pipe);
    }
  }

  private async bufferThenRender(text: string, pipe = process.stdout) {
    this.buffer += text;

    const idx = this.buffer.lastIndexOf("\n#");

    const blockDelims = this.buffer.match(/```/g)?.length;
    if (blockDelims && blockDelims % 2 !== 0) this.unsafeIndex = idx;

    if (idx > this.unsafeIndex) {
      const str = this.buffer.substring(0, idx) ?? "";
      await this.render(str, pipe);
      this.buffer = this.buffer.substring(idx);
      this.unsafeIndex = 0;
    }
  }

  private async render(str: string, pipe = process.stdout) {
    this.moveClearReset(-this.lineCount, pipe);

    const [columns] = pipe.getWindowSize();
    const rendered = await renderMarkdown(str, columns);
    pipe.write(rendered);
    pipe.write("\n");
  }

  async renderStream(stream: AsyncIterableStream<TextStreamPart<ToolSet>>) {
    this.resetCounts();

    for await (const part of stream) {
      switch (part.type) {
        case "reasoning-start": {
          if (supportsColorStderr)
            this.write(
              `${styles.dim.open}${styles.italic.open}`,
              process.stderr,
            );

          this.write("\n", process.stderr);
          break;
        }
        case "reasoning-delta": {
          this.write(part.text, process.stderr);
          if (!this.isPlainErr)
            this.checkContainment(part.text, process.stderr);
          break;
        }
        case "reasoning-end": {
          if (!this.isPlainErr) {
            this.moveClearReset(-this.lineCount - 1, process.stderr);
          } else {
            this.write("\n--- Done Thinking ---\n", process.stderr);
          }

          if (supportsColorStderr)
            this.write(
              `${styles.italic.close}${styles.dim.close}`,
              process.stderr,
            );

          break;
        }
        case "text-start": {
          this.write("\n");
          break;
        }
        case "text-delta": {
          this.write(part.text);
          if (!this.isPlain) {
            this.checkContainment(part.text);
            await this.bufferThenRender(part.text);
          }
          break;
        }
        case "text-end": {
          if (this.buffer) {
            await this.render(this.buffer);
            this.buffer = "";
            this.unsafeIndex = 0;
          } else {
            this.write("\n");
          }
          break;
        }
      }

      while (!this.memOk) await Bun.sleep(10);
    }
  }

  async renderStats(
    usage: PromiseLike<LanguageModelUsage>,
    start: number,
    stop: number,
  ) {
    const { inputTokens, outputTokens } = await usage;
    const context = (inputTokens ?? 0) / 1000;
    const output = (outputTokens ?? 0) / 1000;
    const time = (stop - start) / 1000;
    const template = `Context: ${context.toFixed(1)}k, Output: ${output.toFixed(1)}k, Time: ${time.toFixed(2)}s\n`;
    process.stderr.write(chalkStderr.dim(template));
  }
}