Merge pull request #135 from cyhano/fix/screenshot-media-type

fix: use detected media type for screenshots instead of hard-coded image/png

packages/dsh-plugin-browserskill/src/image.ts

@@ -1,10 +1,11 @@
 /**
- * Screenshot delivery. The canonical tool value stays plain JSON (a PNG file
- * path plus pixel metadata); when the host mounts a durable attachment store
- * AND the calling route declares image input, the PNG bytes are additionally
- * committed through `ctx.attachments` so the render step can attach the image
- * itself to the tool result. Any uncertainty (no store, unknown route,
- * text-only model) falls back to the path-only form.
+ * Screenshot delivery. The canonical tool value stays plain JSON (a screenshot
+ * file path plus pixel metadata); when the host mounts a durable attachment
+ * store AND the calling route declares image input, the capture bytes are
+ * additionally committed through `ctx.attachments` so the render step can
+ * attach the image itself to the tool result. Any uncertainty (no store,
+ * unknown route, text-only model, unrecognized bytes) falls back to the
+ * path-only form.
  */
 
 import type { Context } from "@deepseek-ai/cordis";
@@ -33,6 +34,54 @@ interface LlmLike {
   ): Promise<{ inputModalities?: readonly string[] }>;
 }
 
+/**
+ * Sniff the image media type from magic bytes. The capture path can hand back
+ * JPEG (e.g. `chrome.tabs.captureVisibleTab` on some Chromium/Edge builds
+ * returning JPEG regardless of the requested format), so the declared type can
+ * never be trusted from the file extension alone.
+ * @param data - first bytes of the capture.
+ * @returns the detected media type, or undefined when unrecognized.
+ */
+function sniffImageMediaType(data: Uint8Array): string | undefined {
+  if (data.length >= 4) {
+    // PNG signature: 89 50 4E 47 (0D 0A 1A 0A)
+    if (data[0] === 0x89 && data[1] === 0x50 && data[2] === 0x4e && data[3] === 0x47) {
+      return "image/png";
+    }
+    // JPEG SOI: FF D8 FF
+    if (data[0] === 0xff && data[1] === 0xd8 && data[2] === 0xff) {
+      return "image/jpeg";
+    }
+    // WebP: "RIFF" (52 49 46 46) + "WEBP" at offset 8
+    if (
+      data[0] === 0x52 &&
+      data[1] === 0x49 &&
+      data[2] === 0x46 &&
+      data[3] === 0x46 &&
+      data.length >= 12 &&
+      data[8] === 0x57 &&
+      data[9] === 0x45 &&
+      data[10] === 0x42 &&
+      data[11] === 0x50
+    ) {
+      return "image/webp";
+    }
+    // GIF: "GIF87a" or "GIF89a"
+    if (
+      data[0] === 0x47 &&
+      data[1] === 0x49 &&
+      data[2] === 0x46 &&
+      data[3] === 0x38 &&
+      data.length >= 6 &&
+      (data[4] === 0x37 || data[4] === 0x39) &&
+      data[5] === 0x61
+    ) {
+      return "image/gif";
+    }
+  }
+  return undefined;
+}
+
 /**
  * Commit screenshot bytes to the host attachment store when the composition
  * supports durable images on the current model route.
@@ -46,7 +95,14 @@ export async function trySaveScreenshot(
 ): Promise<ImageAttachmentRef | undefined> {
   const attachments = ctx.get("attachments") as AttachmentLike | undefined;
   if (attachments === undefined) return undefined;
-  if (!attachments.imageLimits.mediaTypes.includes("image/png")) return undefined;
+  // Declare the real media type from the bytes, never a hard-coded "image/png".
+  // Declaring the wrong type trips the store's IMAGE_TYPE_MISMATCH check and
+  // silently degrades to path-only; sniffing first lets JPEG captures (returned
+  // by some Chromium builds for `captureVisibleTab({ format: "png" })`) still
+  // inline correctly.
+  const mediaType = sniffImageMediaType(data);
+  if (mediaType === undefined) return undefined;
+  if (!attachments.imageLimits.mediaTypes.includes(mediaType)) return undefined;
   if (
     data.byteLength >
     Math.min(attachments.imageLimits.maxImageBytes, attachments.imageLimits.maxMessageImageBytes)
@@ -55,9 +111,9 @@ export async function trySaveScreenshot(
   }
   if (!(await isImageCapableRoute(ctx, exec))) return undefined;
   try {
-    return await attachments.saveImage({ data, mediaType: "image/png", name });
+    return await attachments.saveImage({ data, mediaType, name });
   } catch {
-    // A store hiccup must not fail the tool call; the PNG path still works.
+    // A store hiccup must not fail the tool call; the path-only form still works.
     return undefined;
   }
 }

packages/dsh-plugin-browserskill/src/tools.ts

@@ -784,7 +784,11 @@ export function registerTools(deps: ToolDeps): () => void {
               additionalProperties: false,
               properties: {
                 attachmentId: { type: "string", required: true },
-                mediaType: { type: "string", required: true, const: "image/png" },
+                mediaType: {
+                  type: "string",
+                  required: true,
+                  enum: ["image/png", "image/jpeg", "image/webp", "image/gif"],
+                },
                 bytes: { type: "integer", required: true },
                 width: { type: "integer", required: true },
                 height: { type: "integer", required: true },
@@ -854,7 +858,7 @@ export function registerTools(deps: ToolDeps): () => void {
               ? {
                   image: {
                     attachmentId: String(ref.attachmentId),
-                    mediaType: "image/png" as const,
+                    mediaType: ref.mediaType,
                     bytes: ref.bytes,
                     width: ref.width,
                     height: ref.height,

packages/dsh-plugin-browserskill/tests/tools.test.ts

@@ -794,6 +794,15 @@ describe("browser_screenshot", () => {
     return path;
   }
 
+  function jpegFile(): string {
+    const dir = mkdtempSync(join(tmpdir(), "bsk-test-"));
+    const path = join(dir, "shot.png");
+    // A JPEG SOI marker (FF D8 FF) even though the file is named ".png" —
+    // mirrors captureVisibleTab returning JPEG on some Chromium builds.
+    writeFileSync(path, Buffer.from([0xff, 0xd8, 0xff, 0xdb, 0, 1, 2, 3, 4]));
+    return path;
+  }
+
   it("returns the PNG path when no attachment store is mounted", async () => {
     const path = pngFile();
     const { tools, calls } = setup({
@@ -859,6 +868,50 @@ describe("browser_screenshot", () => {
     expect(rendered.map((block) => block.type)).toEqual(["text", "image"]);
   });
 
+  it("declares the real media type when capture bytes are JPEG, not a hard-coded PNG", async () => {
+    const path = jpegFile();
+    const services = {
+      attachments: {
+        imageLimits: {
+          mediaTypes: ["image/png", "image/jpeg", "image/webp", "image/gif"],
+          maxImageBytes: 10_000_000,
+          maxMessageImageBytes: 10_000_000,
+        },
+        saveImage: async (input: { data: Uint8Array; mediaType: string; name?: string }) => ({
+          attachmentId: "att-jpeg",
+          mediaType: input.mediaType,
+          bytes: input.data.byteLength,
+          width: 800,
+          height: 600,
+          name: input.name,
+        }),
+      },
+      llm: {
+        resolveModelInfo: async () => ({ inputModalities: ["text", "image"] }),
+      },
+    };
+    const { tools } = setup(
+      {
+        "session start": START_REPLY("s1"),
+        screenshot: { tab_id: 7, width: 800, height: 600, format: "png", path, byte_size: 9 },
+      },
+      services,
+    );
+    await startSession(tools);
+    const exec = makeExec({
+      agent: {
+        session: { requestHeader: () => ({ config: { provider: "deepseek", model: "vl" } }) },
+        options: {},
+      },
+    });
+    const screenshot = tools.get("browser_screenshot");
+    const value = (await screenshot?.execute({}, exec)) as {
+      image?: { attachmentId: string; mediaType: string };
+    };
+    // The reference media type must reflect the sniffed JPEG bytes, not "image/png".
+    expect(value.image).toMatchObject({ attachmentId: "att-jpeg", mediaType: "image/jpeg" });
+  });
+
   it("stays path-only when the model route is text-only", async () => {
     const path = pngFile();
     const services = {