scripts/probe_token_budgets.mjs

import { chromium } from "@playwright/test";
import { existsSync } from "node:fs";

const baseUrl = process.argv[2] || "http://localhost:5173/?device=wasm";
const budgets = (process.env.TOKEN_BUDGETS || "80,256,2048,8192")
  .split(",")
  .map((value) => value.trim())
  .filter(Boolean);
const executablePath = process.env.CHROMIUM_PATH || (existsSync("/usr/bin/google-chrome") ? "/usr/bin/google-chrome" : "/snap/bin/chromium");

const results = [];

for (const budget of budgets) {
  let browser;
  let page;
  const consoleLines = [];
  const result = { budget, ok: false };

  try {
    browser = await chromium.launch({
      executablePath,
      headless: true,
      args: ["--no-sandbox", "--disable-dev-shm-usage"],
    });
    page = await browser.newPage();
    page.setDefaultTimeout(1200000);
    page.on("console", (message) => {
      consoleLines.push(`${message.type()}: ${message.text()}`);
    });
    page.on("pageerror", (error) => {
      consoleLines.push(`pageerror: ${error.stack || error.message}`);
    });
    page.on("crash", () => {
      consoleLines.push("page crash");
    });

    const started = Date.now();
    await page.goto(baseUrl, { waitUntil: "networkidle" });
    await page.waitForFunction(() => document.querySelector("#status")?.textContent === "Ready");
    await page.fill("#max-new-tokens", budget);
    await page.fill("#temperature", "0");
    await page.selectOption("#gate-device", "wasm");
    await page.click("#confirm-load-model");
    await page.waitForFunction(() => window.__piWebAgent?.modelReady === true, null, {
      timeout: 600000,
    });
    result.loadMs = Date.now() - started;

    const runStarted = Date.now();
    await page.click("#send");
    await page.waitForFunction(() => document.querySelector("#status")?.textContent === "Agent running", null, {
      timeout: 120000,
    });
    await page.waitForFunction(() => document.querySelector("#status")?.textContent === "Ready", null, {
      timeout: 1200000,
    });
    await page.waitForFunction(() => window.__piWebAgent?.transcript?.includes("42"), null, {
      timeout: 120000,
    });
    result.runMs = Date.now() - runStarted;

    const transcript = await page.evaluate(() => window.__piWebAgent?.transcript || "");
    const files = await page.textContent("#files");
    result.modelStatus = await page.textContent("#model-status");
    result.ok = transcript.includes("42") && files?.includes("hello.js");
    result.transcriptChars = transcript.length;
  } catch (error) {
    result.error = error instanceof Error ? error.message : String(error);
  } finally {
    result.warnings = consoleLines.filter((line) => line.startsWith("warning:") || line === "page crash").slice(-3);
    results.push(result);
    await page?.close().catch(() => {});
    await browser?.close().catch(() => {});
  }
}

console.log(JSON.stringify(results, null, 2));
if (results.some((result) => !result.ok)) {
  throw new Error("At least one token budget probe failed.");
}