all-by-skill/skills/gemini-music-web/scripts/run_music_flow.py

#!/usr/bin/env python3
"""Run Gemini music generation flow end-to-end via Playwright CLI."""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
from pathlib import Path


class FlowError(RuntimeError):
    """Raised when a subprocess command in the flow fails."""


def run_command(
    cmd: list[str], *, capture_output: bool = True, check: bool = True
) -> subprocess.CompletedProcess[str]:
    kwargs: dict[str, object] = {"text": True}
    if capture_output:
        kwargs["stdout"] = subprocess.PIPE
        kwargs["stderr"] = subprocess.STDOUT
    proc = subprocess.run(cmd, **kwargs)
    if check and proc.returncode != 0:
        output = proc.stdout if capture_output else ""
        raise FlowError(
            f"Command failed ({proc.returncode}): {' '.join(cmd)}\n{output}"
        )
    return proc


def run_pw(pw_shared: Path, *args: str) -> str:
    proc = run_command([str(pw_shared), *args], capture_output=True)
    return proc.stdout or ""


def is_login_required(pw_shared: Path) -> bool:
    out = run_pw(
        pw_shared,
        "eval",
        (
            "() => {"
            "const hasAccount = !!document.querySelector("
            "'button[aria-label*=\\\"Google 账号\\\"], "
            "button[aria-label*=\\\"Google Account\\\"]'"
            ");"
            "const hasService = !!document.querySelector('a[href*=\\\"ServiceLogin\\\"]');"
            "const hasLoginCtl = Array.from(document.querySelectorAll('a,button'))"
            ".some(el => /登录|Sign in/i.test((el.textContent || '').trim()));"
            "return !hasAccount && (hasService || hasLoginCtl);"
            "}"
        ),
    )
    return bool(re.search(r"(?m)^true$", out))


def enter_music_tool(pw_shared: Path) -> None:
    js = r"""
async (page) => {
const labels = [/创作音乐/, /制作音乐/, /Create music/i, /Music/i];

const tryCardButtons = async () => {
  for (const re of labels) {
    const btn = page.getByRole('button', { name: re }).first();
    if (await btn.count()) {
      try {
        await btn.click({ timeout: 2000 });
        return true;
      } catch (_) {
        // Overlay may intercept pointer. Fall through to menu strategy.
      }
    }
  }
  return false;
};

const tryToolMenu = async () => {
  await page.getByRole('button', { name: '工具', exact: true }).click();
  for (const re of labels) {
    const itemCheck = page.getByRole('menuitemcheckbox', { name: re }).first();
    if (await itemCheck.count()) {
      await itemCheck.click();
      return true;
    }
    const itemPlain = page.getByRole('menuitem', { name: re }).first();
    if (await itemPlain.count()) {
      await itemPlain.click();
      return true;
    }
  }
  return false;
};

let ok = await tryCardButtons();
if (!ok) ok = await tryToolMenu();
if (!ok) ok = await tryToolMenu();
if (!ok) throw new Error('Music tool entry not found');
}
"""
    run_pw(pw_shared, "run-code", js)


def submit_and_download_one(pw_shared: Path, prompt: str) -> None:
    js = f"""
async (page) => {{
const prompt = {json.dumps(prompt)};
const input = page.getByRole('textbox', {{ name: /为 Gemini 输入提示|Enter a prompt/i }}).first();
await input.click();
await input.fill(prompt);
await input.press('Enter');

const stopBtn = page.getByRole('button', {{ name: /停止回答|Stop response/i }}).first();
await stopBtn.waitFor({{ state: 'visible', timeout: 15000 }}).catch(() => {{}});
await stopBtn.waitFor({{ state: 'hidden', timeout: 240000 }});

const downloadBtn = page.getByRole('button', {{ name: /下载音乐作品|Download music/i }}).last();
await downloadBtn.click();

const mp3Item = page.getByRole('menuitem', {{ name: /纯音频|MP3/i }}).first();
if (await mp3Item.count()) {{
  await mp3Item.click();
}} else {{
  const anyItem = page.getByRole('menuitem').first();
  if (await anyItem.count()) await anyItem.click();
}}

await page.waitForTimeout(1200);
}}
"""
    run_pw(pw_shared, "run-code", js)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate music on Gemini web and collect downloaded files."
    )
    parser.add_argument("--prompt", required=True, help="Prompt text for music generation.")
    parser.add_argument(
        "--target", required=True, help="Absolute output directory for collected files."
    )
    parser.add_argument(
        "--count", type=int, default=1, help="Number of tracks to generate. Default: 1."
    )
    parser.add_argument(
        "--no-headed",
        action="store_true",
        help="Run browser without headed mode.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.count < 1:
        print("--count must be a positive integer.", file=sys.stderr)
        return 1

    repo_root = Path(__file__).resolve().parents[3]
    pw_shared = Path(
        os.environ.get("PW_SHARED_WRAPPER", str(repo_root / "tools/pw"))
    ).expanduser()
    collect_script = (Path(__file__).resolve().parent / "collect_downloads.py").resolve()

    if not pw_shared.exists() or not pw_shared.is_file():
        print(f"Shared Playwright wrapper not found: {pw_shared}", file=sys.stderr)
        return 1
    if not os.access(pw_shared, os.X_OK):
        print(f"Shared Playwright wrapper is not executable: {pw_shared}", file=sys.stderr)
        return 1
    if not collect_script.exists():
        print(f"Collector script not found: {collect_script}", file=sys.stderr)
        return 1

    target = Path(args.target).expanduser().resolve()
    target.mkdir(parents=True, exist_ok=True)
    start_ts = time.time()

    try:
        os.environ["PLAYWRIGHT_SHARED_INIT_MODE"] = (
            "headless" if args.no_headed else "headed"
        )
        run_pw(pw_shared, "snapshot")
        run_pw(pw_shared, "goto", "https://gemini.google.com/app")
        run_pw(pw_shared, "snapshot")

        if is_login_required(pw_shared):
            print(
                "Gemini is not logged in. Please log in at https://gemini.google.com/app and rerun.",
                file=sys.stderr,
            )
            return 2

        enter_music_tool(pw_shared)

        for i in range(1, args.count + 1):
            current_prompt = args.prompt
            if args.count > 1:
                current_prompt = (
                    f"{args.prompt}\n"
                    f"变体要求：这是第 {i} / {args.count} 首。保持风格一致，但旋律和节奏细节需要变化。"
                )
            submit_and_download_one(pw_shared, current_prompt)

        collect_cmd = [
            sys.executable,
            str(collect_script),
            "--target",
            str(target),
            "--since",
            str(start_ts),
            "--expected-count",
            str(args.count),
            "--limit",
            str(args.count),
            "--prefix",
            "gemini-music",
            "--prompt",
            args.prompt,
        ]
        proc = run_command(collect_cmd, capture_output=False, check=False)
        return proc.returncode
    except FlowError as exc:
        print(str(exc), file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())