refactor: use context agnostic write to file

Author: Lightning00Blade

src/McpContext.ts

@@ -58,6 +58,20 @@ function getNetworkMultiplierFromString(condition: string | null): number {
   return 1;
 }
 
+function getExtensionFromMimeType(
+  mimeType: 'image/png' | 'image/jpeg' | 'image/webp',
+) {
+  switch (mimeType) {
+    case 'image/png':
+      return 'png';
+    case 'image/jpeg':
+      return 'jpeg';
+    case 'image/webp':
+      return 'webp';
+  }
+  throw new Error(`No mapping for Mime type ${mimeType}.`);
+}
+
 export class McpContext implements Context {
   browser: Browser;
   logger: Debugger;
@@ -346,20 +360,31 @@ export class McpContext implements Context {
       const dir = await fs.mkdtemp(
         path.join(os.tmpdir(), 'chrome-devtools-mcp-'),
       );
-      const ext =
-        mimeType === 'image/png'
-          ? 'png'
-          : mimeType === 'image/jpeg'
-            ? 'jpg'
-            : 'webp';
-      const filename = path.join(dir, `screenshot.${ext}`);
+
+      const filename = path.join(
+        dir,
+        `screenshot.${getExtensionFromMimeType(mimeType)}`,
+      );
       await fs.writeFile(filename, data);
       return {filename};
     } catch (err) {
       this.logger(err);
       throw new Error('Could not save a screenshot to a file');
     }
   }
+  async saveFile(
+    data: Uint8Array<ArrayBufferLike>,
+    filename: string,
+  ): Promise<{filename: string}> {
+    try {
+      const filePath = path.resolve(filename);
+      await fs.writeFile(filePath, data);
+      return {filename};
+    } catch (err) {
+      this.logger(err);
+      throw new Error('Could not save a screenshot to a file');
+    }
+  }
 
   storeTraceRecording(result: TraceResult): void {
     this.#traceResults.push(result);

src/tools/ToolDefinition.ts

@@ -74,6 +74,10 @@ export type Context = Readonly<{
     data: Uint8Array<ArrayBufferLike>,
     mimeType: 'image/png' | 'image/jpeg' | 'image/webp',
   ): Promise<{filename: string}>;
+  saveFile(
+    data: Uint8Array<ArrayBufferLike>,
+    filename: string,
+  ): Promise<{filename: string}>;
   waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
 }>;
 

src/tools/screenshot.ts

@@ -4,8 +4,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import {writeFile} from 'node:fs/promises';
-
 import type {ElementHandle, Page} from 'puppeteer-core';
 import z from 'zod';
 
@@ -85,7 +83,7 @@ export const screenshot = defineTool({
     }
 
     if (request.params.filePath) {
-      await writeFile(request.params.filePath, screenshot);
+      await context.saveFile(screenshot, request.params.filePath);
       response.appendResponseLine(
         `Saved screenshot to ${request.params.filePath}.`,
       );