feat(demo): 新增图片生成与保存功能，优化图片提取逻辑

2026-03-18 02:23:43 +08:00
parent d7f09ff8f7
commit 464327f24b
2 changed files with 130 additions and 39 deletions
--- a/src/demo.js
+++ b/src/demo.js
@@ -15,10 +15,12 @@
 * 所有配置项见 .env，可直接编辑或通过命令行设环境变量。
 */
 import { execSync } from 'node:child_process';
-import { platform } from 'node:os';
+import { platform, homedir } from 'node:os';
+import { writeFileSync, mkdirSync, existsSync } from 'node:fs';
+import { join } from 'node:path';
 import { createGeminiSession, disconnect } from './index.js';

-const prompt = 'Hello Gemini!';
+const prompt = '你好呀~请给我画一张植物大战僵尸的漫画，左侧是豌豆射手，右侧是僵尸。僵尸面对豌豆射手的攻击仓皇逃窜！';

 // ── Demo 专用：杀掉所有 Chromium 系浏览器进程 ──
 function killAllBrowserProcesses() {
@@ -111,16 +113,74 @@ async function main() {
    const probe = await ops.probe();
    console.log('probe:', JSON.stringify(probe, null, 2));

-    // 3. 发送一句话
-    console.log('\n[3] 发送提示词...');
+    // 3. 确保使用 Pro 模型
+    console.log('\n[3] 检查模型...');
+    if (probe.currentModel.toLowerCase() === 'pro') {
+      console.log('[3] 当前已是 Pro 模型，跳过');
+    } else {
+      console.log(`[3] 当前模型: ${probe.currentModel || '未知'}，切换到 Pro...`);
+      const switchResult = await ops.ensureModelPro();
+      if (switchResult.ok) {
+        console.log(`[3] 已切换到 Pro（之前: ${switchResult.previousModel || '未知'}）`);
+      } else {
+        console.warn(`[3] 切换 Pro 失败: ${switchResult.error}，继续使用当前模型`);
+      }
+    }
+
+    // 4. 发送一句话
+    console.log('\n[4] 发送提示词...');
    const result = await ops.sendAndWait(prompt, {
-      timeout: 60_000,
+      timeout: 120_000,
      onPoll(poll) {
        console.log(`  polling... status=${poll.status}`);
      },
    });
    console.log('result:', JSON.stringify(result, null, 2));

+    // 5. 获取最新图片并保存到本地
+    if (result.ok) {
+      console.log('\n[5] 查找最新生成的图片...');
+      await sleep(2000); // 等待图片渲染完毕
+
+      const imgInfo = await ops.getLatestImage();
+      console.log('imgInfo:', JSON.stringify(imgInfo, null, 2));
+
+      if (imgInfo.ok && imgInfo.src) {
+        console.log(`[5] 找到图片 (${imgInfo.width}x${imgInfo.height}, isNew=${imgInfo.isNew})`);
+
+        // 提取 base64 数据
+        console.log(`[5] 提取图片数据 (src=${imgInfo.src})...`);
+        const b64Result = await ops.extractImageBase64(imgInfo.src);
+
+        if (b64Result.ok && b64Result.dataUrl) {
+          // dataUrl 格式: data:image/png;base64,iVBOR...
+          const matches = b64Result.dataUrl.match(/^data:image\/(\w+);base64,(.+)$/);
+          if (matches) {
+            const ext = matches[1] === 'jpeg' ? 'jpg' : matches[1];
+            const base64Data = matches[2];
+            const buffer = Buffer.from(base64Data, 'base64');
+
+            // 保存到 ~/gemini-skill-output/
+            const outputDir = join(homedir(), 'gemini-skill-output');
+            if (!existsSync(outputDir)) {
+              mkdirSync(outputDir, { recursive: true });
+            }
+            const filename = `gemini_${Date.now()}.${ext}`;
+            const filepath = join(outputDir, filename);
+
+            writeFileSync(filepath, buffer);
+            console.log(`[5] ✅ 图片已保存: ${filepath} (${(buffer.length / 1024).toFixed(1)} KB, method=${b64Result.method})`);
+          } else {
+            console.warn('[5] ⚠ dataUrl 格式无法解析');
+          }
+        } else {
+          console.warn(`[5] ⚠ 提取图片数据失败: ${b64Result.error || 'unknown'}`);
+        }
+      } else {
+        console.log('[5] 未找到图片（可能本次回答不含图片）');
+      }
+    }
+
  } catch (err) {
    console.error('Error:', err);
  }
--- a/src/gemini-ops.js
+++ b/src/gemini-ops.js
@@ -465,19 +465,29 @@ export function createOps(page) {
    },

    /**
-     * 提取最新图片的 Base64 数据（Canvas 优先，fetch 兜底）
+     * 提取指定图片的 Base64 数据（Canvas 优先，fetch 兜底）
+     *
+     * @param {string} url - 目标图片的 src URL
+     * @returns {Promise<{ok: boolean, dataUrl?: string, width?: number, height?: number, method?: 'canvas'|'fetch', error?: string}>}
     */
-    async extractImageBase64() {
-      return op.query(() => {
+    async extractImageBase64(url) {
+      if (!url) {
+        console.warn('[extractImageBase64] ❌ 未提供 url 参数');
+        return { ok: false, error: 'missing_url' };
+      }
+      console.log(`[extractImageBase64] 🔍 开始提取, url=${url.slice(0, 120)}...`);
+
+      const canvasResult = await op.query((targetUrl) => {
+        // ── 在页面中根据 url 查找匹配的 img 元素 ──
        const imgs = [...document.querySelectorAll('img.image.loaded')];
-        if (!imgs.length) {
-          return { ok: false, error: 'no_loaded_images' };
+        const img = imgs.find(i => i.src === targetUrl);
+        if (!img) {
+          return { ok: false, error: 'img_not_found_by_url', searched: imgs.length };
        }
-        const img = imgs[imgs.length - 1];
        const w = img.naturalWidth || img.width;
        const h = img.naturalHeight || img.height;

-        // 尝试 Canvas 同步提取
+        // ── 尝试 Canvas 同步提取 ──
        try {
          const canvas = document.createElement('canvas');
          canvas.width = w;
@@ -485,31 +495,52 @@ export function createOps(page) {
          canvas.getContext('2d').drawImage(img, 0, 0);
          const dataUrl = canvas.toDataURL('image/png');
          return { ok: true, dataUrl, width: w, height: h, method: 'canvas' };
-        } catch { /* canvas tainted, fallback */ }
+        } catch (e) {
+          // canvas tainted（跨域图片），记录原因后降级
+          return { ok: false, needFetch: true, src: img.src, width: w, height: h, canvasError: e.message || String(e) };
+        }
+      }, url);

-        // 标记需要 fetch fallback
-        return { ok: false, needFetch: true, src: img.src, width: w, height: h };
-      }).then(async (result) => {
-        if (result.ok || !result.needFetch) return result;
+      if (canvasResult.ok) {
+        console.log(`[extractImageBase64] ✅ Canvas 提取成功 (${canvasResult.width}x${canvasResult.height})`);
+        return canvasResult;
+      }

-        // Fetch fallback: 在页面上下文中异步执行
-        return page.evaluate(async (src, w, h) => {
-          try {
-            const r = await fetch(src);
-            if (!r.ok) throw new Error(`fetch_status_${r.status}`);
-            const blob = await r.blob();
-            return await new Promise((resolve) => {
-              const reader = new FileReader();
-              reader.onloadend = () => resolve({
-                ok: true, dataUrl: reader.result, width: w, height: h, method: 'fetch',
-              });
-              reader.readAsDataURL(blob);
+      if (!canvasResult.needFetch) {
+        // img 元素都没找到，直接返回失败
+        console.warn(`[extractImageBase64] ❌ 页面中未找到匹配的 img 元素 (已扫描 ${canvasResult.searched || 0} 张)`);
+        return canvasResult;
+      }
+
+      // ── Fetch 降级：Canvas 被跨域污染，改用 fetch 读取二进制 ──
+      console.log(`[extractImageBase64] ⚠ Canvas 被污染 (${canvasResult.canvasError})，降级为 fetch...`);
+
+      const fetchResult = await page.evaluate(async (src, w, h) => {
+        try {
+          const r = await fetch(src);
+          if (!r.ok) return { ok: false, error: `fetch_status_${r.status}` };
+          const blob = await r.blob();
+          return await new Promise((resolve) => {
+            const reader = new FileReader();
+            reader.onloadend = () => resolve({
+              ok: true, dataUrl: reader.result, width: w, height: h, method: 'fetch',
            });
-          } catch (err) {
-            return { ok: false, error: 'extract_failed', detail: err.message || String(err) };
-          }
-        }, result.src, result.width, result.height);
-      });
+            reader.onerror = () => resolve({
+              ok: false, error: 'filereader_error',
+            });
+            reader.readAsDataURL(blob);
+          });
+        } catch (err) {
+          return { ok: false, error: 'fetch_failed', detail: err.message || String(err) };
+        }
+      }, canvasResult.src, canvasResult.width, canvasResult.height);
+
+      if (fetchResult.ok) {
+        console.log(`[extractImageBase64] ✅ Fetch 提取成功 (${fetchResult.width}x${fetchResult.height})`);
+      } else {
+        console.warn(`[extractImageBase64] ❌ Fetch 提取失败: ${fetchResult.error}${fetchResult.detail ? ' — ' + fetchResult.detail : ''}`);
+      }
+      return fetchResult;
    },

    /**
@@ -621,12 +652,12 @@ export function createOps(page) {
      await sleep(2000);

      // 4. 获取图片
-      const imgInfo = await this.getLatestImage();
+      let imgInfo = await this.getLatestImage();
      if (!imgInfo.ok) {
        await sleep(3000);
-        const retry = await this.getLatestImage();
-        if (!retry.ok) {
-          return { ok: false, error: 'no_image_found', elapsed: waitResult.elapsed, imgInfo: retry };
+        imgInfo = await this.getLatestImage();
+        if (!imgInfo.ok) {
+          return { ok: false, error: 'no_image_found', elapsed: waitResult.elapsed, imgInfo };
        }
      }

@@ -635,7 +666,7 @@ export function createOps(page) {
        const dlResult = await this.downloadLatestImage();
        return { ok: dlResult.ok, method: 'download', elapsed: waitResult.elapsed, ...dlResult };
      } else {
-        const b64Result = await this.extractImageBase64();
+        const b64Result = await this.extractImageBase64(imgInfo.src);
        return { ok: b64Result.ok, method: b64Result.method, elapsed: waitResult.elapsed, ...b64Result };
      }
    },