feat(core): make screenshot lifecycle format-aware

Author: quanruzhuoxiu

packages/core/src/agent/utils.ts

@@ -23,10 +23,9 @@ import {
 } from '@midscene/shared/env';
 import { generateElementByRect } from '@midscene/shared/extractor';
 import {
-  convertImgBufferToJpeg,
-  createImgBase64ByFormat,
+  canonicalizeScreenshotBase64,
   imageInfoOfBase64,
-  parseBase64,
+  normalizeBase64Image,
   resizeImgBase64,
 } from '@midscene/shared/img';
 import { getDebug } from '@midscene/shared/logger';
@@ -37,28 +36,6 @@ import type { TaskCache } from './task-cache';
 import { debug as cacheDebug } from './task-cache';
 
 const agentDebug = getDebug('agent');
-const screenshotDataUrlPattern = /^data:image\/[a-zA-Z0-9.+-]+;base64,/i;
-
-const inferBase64ImageFormat = (base64Body: string) => {
-  if (base64Body.startsWith('iVBORw0KGgo')) {
-    return 'png';
-  }
-  return 'jpeg';
-};
-
-const normalizeScreenshotBase64 = (screenshotBase64: string) => {
-  const trimmedBase64 = screenshotBase64.trim();
-  if (screenshotDataUrlPattern.test(trimmedBase64)) {
-    return trimmedBase64;
-  }
-
-  const base64Body = trimmedBase64.replace(/\s/g, '');
-  assert(base64Body, 'screenshotBase64 must include image data');
-  return createImgBase64ByFormat(
-    inferBase64ImageFormat(base64Body),
-    base64Body,
-  );
-};
 
 const legacyScrollTypeMap = {
   once: 'singleAction',
@@ -203,23 +180,13 @@ export async function commonContextParser(
       shrunkShotToLogicalRatio,
     };
   } else {
-    // For screenshots that do not need shrinking, convert PNG to JPEG to reduce the image payload in model requests and reports. (Shrunk images are already JPEG.)
-    // This mainly covers Android's default screenshot path, which produces PNG screenshots.
-    // Compared with conversion on Android, centralizing it here means each platform does not need to handle screenshot formats itself, and allows future output formats such as WebP.
-    // Built-in paths that already output JPEG are unaffected, and custom devices that output JPEG will not be compressed again.
-    // The Web platform already outputs JPEG, so it does not enter this branch. Other built-in device platforms run in Node, where Sharp conversion is fast enough that its extra cost is negligible.
-    let outputScreenshotBase64 = screenshotBase64;
-    const { mimeType, body } = parseBase64(screenshotBase64);
-    if (mimeType.toLowerCase() === 'image/png') {
-      const jpegBuffer = await convertImgBufferToJpeg(
-        Buffer.from(body, 'base64'),
-        90,
-      );
-      outputScreenshotBase64 = createImgBase64ByFormat(
-        'jpeg',
-        jpegBuffer.toString('base64'),
-      );
-    }
+    // PNG/raw producers are encoded once as WebP before the ScreenshotItem is
+    // shared by model requests and reports. Native JPEG sources are preserved
+    // to avoid a second lossy encode for MJPEG/HDC frames.
+    const outputScreenshotBase64 = await canonicalizeScreenshotBase64(
+      screenshotBase64,
+      { preserveJpeg: true },
+    );
 
     return {
       shotSize: {
@@ -242,10 +209,13 @@ export async function createScreenshotBoundUIContext(
     screenshotSize?: Size;
   },
 ): Promise<UIContext> {
-  const normalizedScreenshotBase64 =
-    normalizeScreenshotBase64(screenshotBase64);
-  const actualScreenshotSize = await imageInfoOfBase64(
+  const normalizedScreenshotBase64 = normalizeBase64Image(screenshotBase64);
+  const canonicalScreenshotBase64 = await canonicalizeScreenshotBase64(
     normalizedScreenshotBase64,
+    { preserveJpeg: true },
+  );
+  const actualScreenshotSize = await imageInfoOfBase64(
+    canonicalScreenshotBase64,
   );
   if (
     opt.screenshotSize &&
@@ -262,7 +232,7 @@ export async function createScreenshotBoundUIContext(
   }
 
   return {
-    screenshot: ScreenshotItem.create(normalizedScreenshotBase64, Date.now()),
+    screenshot: ScreenshotItem.create(canonicalScreenshotBase64, Date.now()),
     shotSize: actualScreenshotSize,
     shrunkShotToLogicalRatio: 1,
     _isFrozen: true,

packages/core/src/dump/report-action-dump.ts

@@ -6,14 +6,18 @@ import {
   writeFileSync,
 } from 'node:fs';
 import { join } from 'node:path';
+import {
+  screenshotImageExtension,
+  screenshotImageFormatFromMimeType,
+} from '@midscene/shared/img/image-format';
 import { ScreenshotItem } from '../screenshot-item';
 import type {
   ExecutionTask,
   IExecutionDump,
   IReportActionDump,
 } from '../types';
 import { restoreImageReferences } from './screenshot-restoration';
-import { ScreenshotStore } from './screenshot-store';
+import { type ScreenshotRef, ScreenshotStore } from './screenshot-store';
 
 /**
  * Replacer function for JSON serialization that handles Page, Browser objects and ScreenshotItem
@@ -253,10 +257,10 @@ export class ReportActionDump implements IReportActionDump {
   }
 
   /**
-   * Serialize the dump to files with screenshots as separate PNG files.
+   * Serialize the dump to files with screenshots as separate image files.
    * Creates:
    * - {basePath} - dump JSON with { $screenshot: id } references
-   * - {basePath}.screenshots/ - PNG files
+   * - {basePath}.screenshots/ - screenshot image files
    *
    * @param basePath - Base path for the dump file
    */
@@ -296,14 +300,21 @@ export class ReportActionDump implements IReportActionDump {
     const dumpString = readFileSync(basePath, 'utf-8');
     const screenshotsDir = `${basePath}.screenshots`;
 
-    const loadFromExecutionScreenshotDir = (id: string, mimeType: string) => {
-      const ext = mimeType === 'image/jpeg' ? 'jpeg' : 'png';
+    const loadFromExecutionScreenshotDir = (
+      id: string,
+      mimeType: ScreenshotRef['mimeType'],
+    ) => {
+      const format = screenshotImageFormatFromMimeType(mimeType);
+      if (!format) {
+        throw new Error(`Unsupported screenshot mime type: ${mimeType}`);
+      }
+      const ext = screenshotImageExtension(format);
       const filePath = join(screenshotsDir, `${id}.${ext}`);
       if (!existsSync(filePath)) {
         return '';
       }
       const data = readFileSync(filePath);
-      return `data:image/${ext};base64,${data.toString('base64')}`;
+      return `data:${mimeType};base64,${data.toString('base64')}`;
     };
 
     // Restore image references

packages/core/src/dump/screenshot-store.ts

@@ -1,14 +1,20 @@
 import { existsSync, mkdirSync, readFileSync, rmSync } from 'node:fs';
 import { writeFile as writeFileAsync } from 'node:fs/promises';
 import { dirname, isAbsolute, join } from 'node:path';
+import {
+  type ScreenshotImageMimeType,
+  isScreenshotImageMimeType,
+  screenshotImageExtension,
+  screenshotImageFormatFromMimeType,
+} from '@midscene/shared/img/image-format';
 import type { ScreenshotItem } from '../screenshot-item';
 import { extractImageByIdSync } from './html-utils';
 
 export interface ScreenshotRef {
   type: 'midscene_screenshot_ref';
   id: string;
   capturedAt: number;
-  mimeType: 'image/png' | 'image/jpeg';
+  mimeType: ScreenshotImageMimeType;
   storage: 'inline' | 'file';
   path?: string;
 }
@@ -22,7 +28,7 @@ export function normalizeScreenshotRef(value: unknown): ScreenshotRef | null {
     typeof record.id === 'string' &&
     typeof record.capturedAt === 'number' &&
     (record.storage === 'inline' || record.storage === 'file') &&
-    (record.mimeType === 'image/png' || record.mimeType === 'image/jpeg')
+    isScreenshotImageMimeType(record.mimeType)
   ) {
     if (record.storage === 'file' && typeof record.path !== 'string') {
       return null;
@@ -48,7 +54,11 @@ type ResolvedScreenshotSource =
     };
 
 function extensionByMimeType(mimeType: ScreenshotRef['mimeType']): string {
-  return mimeType === 'image/jpeg' ? 'jpeg' : 'png';
+  const format = screenshotImageFormatFromMimeType(mimeType);
+  if (!format) {
+    throw new Error(`Unsupported screenshot mime type: ${mimeType}`);
+  }
+  return screenshotImageExtension(format);
 }
 
 export function resolveScreenshotSource(

packages/core/src/screenshot-item.ts

@@ -1,4 +1,12 @@
 import { readFileSync } from 'node:fs';
+import {
+  type ScreenshotImageFormat,
+  type ScreenshotImageMimeType,
+  inferScreenshotImageFormatFromBase64,
+  screenshotImageExtension,
+  screenshotImageFormatFromMimeType,
+  screenshotImageMimeType,
+} from '@midscene/shared/img/image-format';
 import { uuid } from '@midscene/shared/utils';
 import { extractImageByIdSync } from './dump/html-utils';
 import {
@@ -13,13 +21,32 @@ import {
  */
 export type ScreenshotSerializeFormat = ScreenshotRef;
 
+const BASE64_SEPARATOR = ';base64,';
+
 /**
- * Detect image format from base64 data URI prefix.
+ * Detect image format from a data URI or raw base64 body.
  */
-function detectFormat(base64: string): 'png' | 'jpeg' {
-  if (base64.startsWith('data:image/jpeg')) return 'jpeg';
-  if (base64.startsWith('data:image/jpg')) return 'jpeg';
-  return 'png';
+function detectFormat(base64: string): ScreenshotImageFormat {
+  const separatorIndex = base64.indexOf(BASE64_SEPARATOR);
+  const mimeType =
+    separatorIndex === -1 ? undefined : base64.slice(5, separatorIndex);
+  const detectedFormat =
+    separatorIndex === -1
+      ? inferScreenshotImageFormatFromBase64(base64)
+      : screenshotImageFormatFromMimeType(mimeType);
+
+  // Before WebP support, every non-JPEG value used PNG metadata. Preserve that
+  // behavior for temporary and test placeholders while detecting valid images.
+  return detectedFormat ?? 'png';
+}
+
+function rawBase64Body(base64: string): string {
+  const separatorIndex = base64.indexOf(BASE64_SEPARATOR);
+  const body =
+    separatorIndex === -1
+      ? base64
+      : base64.slice(separatorIndex + BASE64_SEPARATOR.length);
+  return body.replace(/\s/g, '');
 }
 
 /**
@@ -35,7 +62,7 @@ function detectFormat(base64: string): 'png' | 'jpeg' {
 export class ScreenshotItem {
   private _id: string;
   private _base64: string | null;
-  private _format: 'png' | 'jpeg';
+  private _format: ScreenshotImageFormat;
   private _capturedAt: number;
   private _serializedRef: ScreenshotRef | null = null;
   private _persistedPath: string | null = null;
@@ -57,14 +84,19 @@ export class ScreenshotItem {
     return this._id;
   }
 
-  /** Get the image format (png or jpeg) */
-  get format(): 'png' | 'jpeg' {
+  /** Get the image format (PNG, JPEG, or WebP). */
+  get format(): ScreenshotImageFormat {
     return this._format;
   }
 
   /** Get the file extension for this screenshot */
-  get extension(): string {
-    return this._format === 'jpeg' ? 'jpeg' : 'png';
+  get extension(): ScreenshotImageFormat {
+    return screenshotImageExtension(this._format);
+  }
+
+  /** Get the MIME type for this screenshot. */
+  get mimeType(): ScreenshotImageMimeType {
+    return screenshotImageMimeType(this._format);
   }
 
   /** Get screenshot capture timestamp in milliseconds */
@@ -83,7 +115,7 @@ export class ScreenshotItem {
         throw new Error(`Screenshot ${this._id}: file recovery path missing`);
       }
       const buffer = readFileSync(this._persistedPath);
-      return `data:image/${this._format};base64,${buffer.toString('base64')}`;
+      return `data:${this.mimeType};base64,${buffer.toString('base64')}`;
     };
 
     const loadFromInline = (): string => {
@@ -176,7 +208,7 @@ export class ScreenshotItem {
         type: 'midscene_screenshot_ref',
         id: this._id,
         capturedAt: this._capturedAt,
-        mimeType: this._format === 'jpeg' ? 'image/jpeg' : 'image/png',
+        mimeType: this.mimeType,
         storage: 'inline',
       }
     );
@@ -195,7 +227,7 @@ export class ScreenshotItem {
       type: 'midscene_screenshot_ref',
       id: this._id,
       capturedAt: this._capturedAt,
-      mimeType: this._format === 'jpeg' ? 'image/jpeg' : 'image/png',
+      mimeType: this.mimeType,
       storage,
     };
     if (storage === 'file') {
@@ -213,6 +245,6 @@ export class ScreenshotItem {
    * Useful for writing raw binary data to files.
    */
   get rawBase64(): string {
-    return this.base64.replace(/^data:image\/(png|jpeg|jpg);base64,/, '');
+    return rawBase64Body(this.base64);
   }
 }

packages/core/src/types.ts

@@ -489,7 +489,7 @@ export interface ExecutionRecorderItem {
 
 export interface RecordToReportScreenshot {
   /**
-   * PNG/JPEG data URI, or raw PNG base64 body.
+   * PNG/JPEG/WebP data URI, or raw PNG/WebP base64 body.
    */
   base64: string;
   description?: string;
@@ -740,7 +740,7 @@ export type ExecutionTaskPlanningLocate =
 /*
 How a report file stores screenshots:
 - `inline`: base64 image script tags embedded in the single HTML file
-- `directory`: external PNG files under a sibling `screenshots/` dir
+- `directory`: external image files under a sibling `screenshots/` dir
 */
 export type ScreenshotMode = 'inline' | 'directory';
 
@@ -916,7 +916,7 @@ export interface AgentOpt {
    * Use directory-based report format with separate image files.
    *
    * When enabled:
-   * - Screenshots are saved as PNG files in a `screenshots/` subdirectory
+   * - Screenshots retain their image format in a `screenshots/` subdirectory
    * - Report is generated as `index.html` with relative image paths
    * - Reduces memory usage and report file size
    *

packages/core/tests/unit-test/agent-describe-element.test.ts

@@ -424,7 +424,7 @@ describe('element describer utils', () => {
 
     const describeContext = describe.mock.calls[0][2]?.context;
     expect(describeContext?.screenshot.base64).toMatch(
-      /^data:image\/png;base64,/,
+      /^data:image\/webp;base64,UklGR/,
     );
     expect(describeContext?.shotSize).toEqual(fixtureScreenshotSize);
 

packages/core/tests/unit-test/agent-dump-update.test.ts

@@ -255,7 +255,7 @@ describe('Agent dump update screenshot serialization', () => {
         screenshots: [{ base64: 'data:image/svg+xml;base64,custom' }],
       }),
     ).rejects.toThrow(
-      'recordToReport: screenshot #1 base64 must be a PNG/JPEG data URI or raw PNG base64 string',
+      'recordToReport: screenshot #1 base64 must be a PNG/JPEG/WebP data URI or raw PNG/WebP base64 string',
     );
 
     expect(screenshotBase64).not.toHaveBeenCalled();

packages/core/tests/unit-test/common-context-parser-orientation.test.ts

@@ -4,13 +4,8 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 // Mock imageInfoOfBase64 to control screenshot dimensions
 vi.mock('@midscene/shared/img', () => ({
-  convertImgBufferToJpeg: vi.fn(),
-  createImgBase64ByFormat: vi.fn(),
+  canonicalizeScreenshotBase64: vi.fn().mockResolvedValue('mock-base64-data'),
   imageInfoOfBase64: vi.fn(),
-  parseBase64: vi.fn(() => ({
-    mimeType: 'image/jpeg',
-    body: 'mock-base64-data',
-  })),
   resizeImgBase64: vi.fn().mockResolvedValue('mock-resized-base64-data'),
 }));