refactor: use context agnostic write to file (#250)

Similar to the write to temporary file we should hide the implementation
behind the Context

Author: Lightning00Blade

src/McpContext.ts

@@ -58,6 +58,18 @@ function getNetworkMultiplierFromString(condition: string | null): number {
   return 1;
 }
 
+function getExtensionFromMimeType(mimeType: string) {
+  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,18 +358,29 @@ 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');
+      throw new Error('Could not save a screenshot to a file', {cause: err});
+    }
+  }
+  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', {cause: err});
     }
   }
 

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,10 +83,8 @@ export const screenshot = defineTool({
     }
 
     if (request.params.filePath) {
-      await writeFile(request.params.filePath, screenshot);
-      response.appendResponseLine(
-        `Saved screenshot to ${request.params.filePath}.`,
-      );
+      const file = await context.saveFile(screenshot, request.params.filePath);
+      response.appendResponseLine(`Saved screenshot to ${file.filename}.`);
     } else if (screenshot.length >= 2_000_000) {
       const {filename} = await context.saveTemporaryFile(
         screenshot,