ai-chatbot/scripts/nanobanana-browser-poc.ts

import { chromium, type ConsoleMessage, type Page, type Request, type Response } from "@playwright/test";
import { mkdir, writeFile } from "node:fs/promises";
import path from "node:path";
import process from "node:process";

type Mode = "remove-background" | "upscale";

type Options = {
  autoFill: boolean;
  autoSubmit: boolean;
  closeOnSuccess: boolean;
  headless: boolean;
  imageUrl?: string;
  locale: string;
  mode: Mode;
  outDir: string;
  timeoutMs: number;
};

type EventKind =
  | "console"
  | "download"
  | "image-safety"
  | "pageerror"
  | "presigned-urls"
  | "remove-background"
  | "result"
  | "upscale";

type MatchedEvent = {
  kind: EventKind;
  note?: string;
  url: string;
};

type CaptureSummary = {
  artifactDir: string;
  consoleCount: number;
  downloads: string[];
  errorCount: number;
  eventCount: number;
  finalImageUrl?: string;
  mode: Mode;
  pageUrl: string;
  requestCount: number;
  responseCount: number;
  successDetected: boolean;
};

const DEFAULT_IMAGE_URL =
  "https://upload.wikimedia.org/wikipedia/commons/thumb/a/a9/Example.jpg/240px-Example.jpg";

const MODE_CONFIG: Record<
  Mode,
  {
    actionLabels: string[];
    pagePath: string;
  }
> = {
  "remove-background": {
    actionLabels: ["移除背景", "Remove Background"],
    pagePath: "remove-background",
  },
  upscale: {
    actionLabels: ["放大图片", "Upscale Image"],
    pagePath: "ai-image-upscaler",
  },
};

function printHelp() {
  console.log(`Nano Banana browser POC

Usage:
  pnpm poc:nanobanana --mode <remove-background|upscale> [options]

Options:
  --mode <value>             Required. One of: remove-background, upscale
  --image-url <value>        Public image URL to prefill. Defaults to a sample Wikimedia image.
  --locale <value>           Site locale. Defaults to zh.
  --out-dir <value>          Output directory. Defaults to playwright/nanobanana-poc/<timestamp>.
  --headless                 Run headless. Default is visible browser.
  --no-auto-fill             Do not try to prefill the image URL field.
  --auto-submit              Try to click the primary action button after prefilling.
  --close-on-success         Close the browser once a successful result is detected.
  --timeout-ms <value>       Session timeout in milliseconds. Defaults to 600000.
  --help                     Show this message.

Notes:
  - This POC does not bypass Turnstile, login, or rate limits.
  - Complete any CAPTCHA or login step manually in the opened browser.
  - Close the browser window when you are done if --close-on-success is not used.
`);
}

function timestampLabel(date = new Date()) {
  const parts = [
    date.getFullYear(),
    `${date.getMonth() + 1}`.padStart(2, "0"),
    `${date.getDate()}`.padStart(2, "0"),
    `${date.getHours()}`.padStart(2, "0"),
    `${date.getMinutes()}`.padStart(2, "0"),
    `${date.getSeconds()}`.padStart(2, "0"),
  ];
  return `${parts[0]}${parts[1]}${parts[2]}-${parts[3]}${parts[4]}${parts[5]}`;
}

function sanitizeFilePart(value: string) {
  return value.replace(/[^a-zA-Z0-9._-]+/g, "-").slice(0, 80) || "item";
}

function getFlagValue(args: string[], flag: string) {
  const index = args.indexOf(flag);
  if (index === -1) {
    return undefined;
  }

  const value = args[index + 1];
  if (!value || value.startsWith("--")) {
    throw new Error(`Missing value for ${flag}`);
  }

  return value;
}

function hasFlag(args: string[], flag: string) {
  return args.includes(flag);
}

function parseArgs(argv: string[]): Options {
  if (hasFlag(argv, "--help")) {
    printHelp();
    process.exit(0);
  }

  const mode = getFlagValue(argv, "--mode");
  if (mode !== "remove-background" && mode !== "upscale") {
    throw new Error("You must pass --mode remove-background or --mode upscale");
  }

  const locale = getFlagValue(argv, "--locale") ?? "zh";
  const imageUrl = getFlagValue(argv, "--image-url") ?? DEFAULT_IMAGE_URL;
  const timeoutMsRaw = getFlagValue(argv, "--timeout-ms");
  const timeoutMs = timeoutMsRaw ? Number.parseInt(timeoutMsRaw, 10) : 10 * 60 * 1000;

  if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
    throw new Error("--timeout-ms must be a positive integer");
  }

  const outDir =
    getFlagValue(argv, "--out-dir") ??
    path.join(process.cwd(), "playwright", "nanobanana-poc", `${mode}-${timestampLabel()}`);

  return {
    autoFill: !hasFlag(argv, "--no-auto-fill"),
    autoSubmit: hasFlag(argv, "--auto-submit"),
    closeOnSuccess: hasFlag(argv, "--close-on-success"),
    headless: hasFlag(argv, "--headless"),
    imageUrl,
    locale,
    mode,
    outDir,
    timeoutMs,
  };
}

function matchInterestingEvent(url: string, mode: Mode): MatchedEvent | null {
  if (url.includes("/api/img/image-safety")) {
    return { kind: "image-safety", url };
  }

  if (url.includes("/api/presigned-urls")) {
    return { kind: "presigned-urls", url };
  }

  if (url.includes("/api/img/result/")) {
    return { kind: "result", url };
  }

  if (mode === "remove-background" && url.includes("/api/img/rmbg")) {
    return { kind: "remove-background", url };
  }

  if (mode === "upscale" && url.includes("/api/img/upscale")) {
    return { kind: "upscale", url };
  }

  return null;
}

function safeJsonParse(value: string) {
  try {
    return JSON.parse(value);
  } catch {
    return undefined;
  }
}

function getTargetUrl(mode: Mode, locale: string) {
  return `https://nanobanana.im/${locale}/${MODE_CONFIG[mode].pagePath}`;
}

async function ensureDir(dirPath: string) {
  await mkdir(dirPath, { recursive: true });
}

async function writeJson(filePath: string, value: unknown) {
  await writeFile(filePath, `${JSON.stringify(value, null, 2)}\n`, "utf-8");
}

async function saveScreenshot(page: Page, outDir: string, name: string) {
  const screenshotPath = path.join(outDir, `${name}.png`);
  await page.screenshot({ fullPage: true, path: screenshotPath });
  return screenshotPath;
}

async function tryClickByLabels(page: Page, role: "button" | "tab", labels: string[]) {
  for (const label of labels) {
    const locator = page.getByRole(role, { name: label, exact: false });
    if ((await locator.count()) >= 1) {
      await locator.first().click();
      return label;
    }
  }

  return undefined;
}

async function tryFillImageUrl(page: Page, imageUrl: string) {
  const candidates = [
    page.getByPlaceholder("请输入图片链接...", { exact: false }),
    page.getByPlaceholder("Enter image URL...", { exact: false }),
    page.locator('input[id="image-url"]'),
    page.locator('input[type="url"]'),
    page.locator('input[placeholder*="URL" i]'),
    page.locator('input[placeholder*="url" i]'),
    page.getByRole("textbox", { name: "请输入图片链接...", exact: false }),
    page.getByRole("textbox", { name: "Enter image URL...", exact: false }),
  ];

  for (const locator of candidates) {
    if ((await locator.count()) >= 1) {
      await locator.first().fill(imageUrl);
      return true;
    }
  }

  return false;
}

async function tryPrepareForm(page: Page, options: Options, outDir: string) {
  if (!options.autoFill || !options.imageUrl) {
    return {
      imageUrlFilled: false,
      submitClicked: false,
      urlTabClicked: false,
    };
  }

  const urlTabClicked = Boolean(
    await tryClickByLabels(page, "tab", [
      "图片链接",
      "Image URL",
      "Image Url",
      "URL",
      "Link",
    ])
  );

  const imageUrlFilled = await tryFillImageUrl(page, options.imageUrl);

  if (imageUrlFilled) {
    await saveScreenshot(page, outDir, "02-prefilled");
  }

  let submitClicked = false;
  if (options.autoSubmit) {
    const clickedLabel = await tryClickByLabels(page, "button", MODE_CONFIG[options.mode].actionLabels);
    submitClicked = Boolean(clickedLabel);
    if (submitClicked) {
      await saveScreenshot(page, outDir, "03-after-submit-click");
    }
  }

  return {
    imageUrlFilled,
    submitClicked,
    urlTabClicked,
  };
}

function describeConsoleMessage(message: ConsoleMessage) {
  return {
    location: message.location(),
    text: message.text(),
    type: message.type(),
  };
}

async function run() {
  const options = parseArgs(process.argv.slice(2));
  const requestsDir = path.join(options.outDir, "requests");
  const responsesDir = path.join(options.outDir, "responses");
  const logsDir = path.join(options.outDir, "logs");
  const downloadsDir = path.join(options.outDir, "downloads");

  await Promise.all([ensureDir(options.outDir), ensureDir(requestsDir), ensureDir(responsesDir), ensureDir(logsDir), ensureDir(downloadsDir)]);

  const browser = await chromium.launch({
    headless: options.headless,
  });
  const context = await browser.newContext({
    acceptDownloads: true,
    viewport: {
      height: 1000,
      width: 1440,
    },
  });
  const page = await context.newPage();

  const pendingWrites: Promise<unknown>[] = [];
  const downloads: string[] = [];
  const summary: CaptureSummary = {
    artifactDir: options.outDir,
    consoleCount: 0,
    downloads,
    errorCount: 0,
    eventCount: 0,
    mode: options.mode,
    pageUrl: getTargetUrl(options.mode, options.locale),
    requestCount: 0,
    responseCount: 0,
    successDetected: false,
  };

  let sequence = 0;
  let finishRequested = false;
  let resolveFinished: (() => void) | undefined;
  const finished = new Promise<void>((resolve) => {
    resolveFinished = resolve;
  });

  const requestFinish = () => {
    if (finishRequested) {
      return;
    }
    finishRequested = true;
    resolveFinished?.();
  };

  page.on("console", (message) => {
    summary.consoleCount += 1;
    const record = {
      ...describeConsoleMessage(message),
      timestamp: new Date().toISOString(),
    };
    const fileName = `${String(++sequence).padStart(3, "0")}-console.json`;
    pendingWrites.push(writeJson(path.join(logsDir, fileName), record));
  });

  page.on("pageerror", (error) => {
    summary.errorCount += 1;
    const fileName = `${String(++sequence).padStart(3, "0")}-pageerror.json`;
    pendingWrites.push(
      writeJson(path.join(logsDir, fileName), {
        message: error.message,
        name: error.name,
        stack: error.stack,
        timestamp: new Date().toISOString(),
      })
    );
  });

  page.on("download", async (download) => {
    const suggestedName = sanitizeFilePart(download.suggestedFilename());
    const downloadPath = path.join(downloadsDir, suggestedName);
    downloads.push(downloadPath);
    await download.saveAs(downloadPath);
    const fileName = `${String(++sequence).padStart(3, "0")}-download.json`;
    pendingWrites.push(
      writeJson(path.join(logsDir, fileName), {
        path: downloadPath,
        suggestedFilename: download.suggestedFilename(),
        timestamp: new Date().toISOString(),
      })
    );
  });

  page.on("request", (request: Request) => {
    const matched = matchInterestingEvent(request.url(), options.mode);
    if (!matched) {
      return;
    }

    summary.eventCount += 1;
    summary.requestCount += 1;
    const fileName = `${String(++sequence).padStart(3, "0")}-${matched.kind}-request.json`;
    pendingWrites.push(
      writeJson(path.join(requestsDir, fileName), {
        headers: request.headers(),
        kind: matched.kind,
        method: request.method(),
        postData: request.postData(),
        resourceType: request.resourceType(),
        timestamp: new Date().toISOString(),
        url: request.url(),
      })
    );
  });

  page.on("response", async (response: Response) => {
    const matched = matchInterestingEvent(response.url(), options.mode);
    if (!matched) {
      return;
    }

    summary.eventCount += 1;
    summary.responseCount += 1;

    const bodyText = await response.text().catch(() => undefined);
    const bodyJson = bodyText ? safeJsonParse(bodyText) : undefined;

    if (matched.kind === "remove-background" && bodyJson?.data && typeof bodyJson.data === "string") {
      summary.finalImageUrl = bodyJson.data;
      summary.successDetected = true;
      if (options.closeOnSuccess) {
        requestFinish();
      }
    }

    if (matched.kind === "result") {
      const candidateUrl =
        bodyJson?.data?.imgAfterSrc ?? bodyJson?.imgAfterSrc;
      if (typeof candidateUrl === "string" && candidateUrl.length > 0) {
        summary.finalImageUrl = candidateUrl;
        summary.successDetected = true;
        if (options.closeOnSuccess) {
          requestFinish();
        }
      }
    }

    const fileName = `${String(++sequence).padStart(3, "0")}-${matched.kind}-response.json`;
    pendingWrites.push(
      writeJson(path.join(responsesDir, fileName), {
        bodyJson,
        bodyText,
        headers: await response.allHeaders(),
        kind: matched.kind,
        ok: response.ok(),
        status: response.status(),
        statusText: response.statusText(),
        timestamp: new Date().toISOString(),
        url: response.url(),
      })
    );
  });

  const timeout = setTimeout(() => {
    console.log(`Session timed out after ${options.timeoutMs}ms`);
    requestFinish();
  }, options.timeoutMs);

  page.on("close", () => {
    requestFinish();
  });

  const pageUrl = getTargetUrl(options.mode, options.locale);
  console.log(`Opening ${pageUrl}`);
  console.log(`Artifacts will be written to ${options.outDir}`);
  console.log("Manual steps are still expected for Turnstile, login, and any final confirmation.");

  await page.goto(pageUrl, { waitUntil: "domcontentloaded" });
  await saveScreenshot(page, options.outDir, "01-loaded");
  await writeJson(path.join(options.outDir, "storage-state-initial.json"), await context.storageState());

  const formPrep = await tryPrepareForm(page, options, options.outDir);
  await writeJson(path.join(options.outDir, "session-plan.json"), {
    ...formPrep,
    mode: options.mode,
    options: {
      autoFill: options.autoFill,
      autoSubmit: options.autoSubmit,
      closeOnSuccess: options.closeOnSuccess,
      headless: options.headless,
      imageUrl: options.imageUrl,
      locale: options.locale,
      timeoutMs: options.timeoutMs,
    },
    pageUrl,
    startedAt: new Date().toISOString(),
  });

  console.log("Browser is ready.");
  if (!formPrep.imageUrlFilled) {
    console.log("The script did not prefill the image URL. Fill the form manually if needed.");
  }
  if (options.autoSubmit && !formPrep.submitClicked) {
    console.log("The script did not find the primary action button. Click it manually.");
  }
  console.log("Close the browser window when you are done, or wait for timeout.");

  process.on("SIGINT", () => {
    console.log("Received SIGINT, finishing capture...");
    requestFinish();
  });

  await finished;
  clearTimeout(timeout);

  await saveScreenshot(page, options.outDir, "99-final").catch(() => undefined);
  await writeJson(path.join(options.outDir, "storage-state-final.json"), await context.storageState());
  await Promise.allSettled(pendingWrites);

  await writeJson(path.join(options.outDir, "summary.json"), {
    ...summary,
    finishedAt: new Date().toISOString(),
  });

  await context.close();
  await browser.close();

  console.log("");
  console.log("POC finished.");
  console.log(JSON.stringify(summary, null, 2));
}

run().catch((error) => {
  console.error(error);
  process.exit(1);
});