From c9c82470e8b922a94abcbcd96bc5d2828ac0449a Mon Sep 17 00:00:00 2001
From: Ryan Malloy <ryan@supported.systems>
Date: Thu, 4 Sep 2025 14:46:35 -0600
Subject: [PATCH] fix: prevent large screenshots from being sent to API while
 preserving file save
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Block images exceeding 8000px from being included in MCP responses
- Add clear warning messages when images are too large for API
- Always save screenshots to file regardless of size
- Prevents conversation history issues with oversized fullPage screenshots
- Add test script for verifying large image protection

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 src/tools/screenshot.ts  | 29 ++++++++++++++----
 test-large-screenshot.js | 64 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 5 deletions(-)
 create mode 100755 test-large-screenshot.js

diff --git a/src/tools/screenshot.ts b/src/tools/screenshot.ts
index 341e037..3e142a7 100644
--- a/src/tools/screenshot.ts
+++ b/src/tools/screenshot.ts
@@ -62,7 +62,7 @@ const screenshotSchema = z.object({
   filename: z.string().optional().describe('File name to save the screenshot to. Defaults to `page-{timestamp}.{png|jpeg}` if not specified.'),
   element: z.string().optional().describe('Human-readable element description used to obtain permission to screenshot the element. If not provided, the screenshot will be taken of viewport. If element is provided, ref must be provided too.'),
   ref: z.string().optional().describe('Exact target element reference from the page snapshot. If not provided, the screenshot will be taken of viewport. If ref is provided, element must be provided too.'),
-  fullPage: z.boolean().optional().describe('When true, takes a screenshot of the full scrollable page, instead of the currently visible viewport. Cannot be used with element screenshots.'),
+  fullPage: z.boolean().optional().describe('When true, takes a screenshot of the full scrollable page, instead of the currently visible viewport. Cannot be used with element screenshots. WARNING: Full page screenshots may exceed API size limits on long pages.'),
   allowLargeImages: z.boolean().optional().describe('Allow images with dimensions exceeding 8000 pixels (API limit). Default false - will error if image is too large to prevent API failures.'),
 }).refine(data => {
   return !!data.element === !!data.ref;
@@ -160,10 +160,29 @@ const screenshot = defineTabTool({
     }
 
     response.addResult(resultMessage);
-    response.addImage({
-      contentType: fileType === 'png' ? 'image/png' : 'image/jpeg',
-      data: buffer
-    });
+    
+    // Only add image to response if dimensions are safe or explicitly allowed
+    let addImageToResponse = true;
+    if (!params.allowLargeImages) {
+      try {
+        const { width, height } = getImageDimensions(buffer);
+        const maxDimension = 8000;
+        if (width > maxDimension || height > maxDimension) {
+          addImageToResponse = false;
+        }
+      } catch (dimensionError) {
+        // If we can't parse dimensions, continue and add the image
+      }
+    }
+    
+    if (addImageToResponse) {
+      response.addImage({
+        contentType: fileType === 'png' ? 'image/png' : 'image/jpeg',
+        data: buffer
+      });
+    } else {
+      response.addResult(`\n\n🚫 **Image not included in response**: Screenshot exceeds API size limits (8000px). Image saved to file only.`);
+    }
   }
 });
 
diff --git a/test-large-screenshot.js b/test-large-screenshot.js
new file mode 100755
index 0000000..73fd354
--- /dev/null
+++ b/test-large-screenshot.js
@@ -0,0 +1,64 @@
+#!/usr/bin/env node
+
+/**
+ * Test script to verify large screenshot handling
+ * Creates a very tall page and tests fullPage screenshot protection
+ */
+
+const { execSync } = require('child_process');
+const fs = require('fs');
+const path = require('path');
+
+async function testLargeScreenshot() {
+  console.log('🧪 Testing large screenshot protection...');
+  
+  // Create a simple HTML page that will be very tall
+  const testHtml = `
+<!DOCTYPE html>
+<html>
+<head>
+  <title>Large Page Test</title>
+  <style>
+    .tall-content {
+      height: 10000px;
+      background: linear-gradient(to bottom, #ff0000, #00ff00, #0000ff);
+      display: flex;
+      align-items: center;
+      justify-content: center;
+      font-size: 48px;
+      color: white;
+      text-shadow: 2px 2px 4px rgba(0,0,0,0.5);
+    }
+  </style>
+</head>
+<body>
+  <div class="tall-content">
+    This is a very tall page (10000px height)<br>
+    Should trigger large image protection
+  </div>
+</body>
+</html>
+  `;
+  
+  const testFile = path.join(__dirname, 'test-large-page.html');
+  fs.writeFileSync(testFile, testHtml);
+  
+  console.log(`📄 Created test file: ${testFile}`);
+  console.log('🔧 This test requires manual verification with an MCP client');
+  console.log('');
+  console.log('To test:');
+  console.log('1. Start MCP server: npm run build && node lib/index.js');
+  console.log(`2. Navigate to: file://${testFile}`);
+  console.log('3. Try: browser_take_screenshot {"fullPage": true}');
+  console.log('4. Verify: Image saved to file but NOT included in response');
+  console.log('5. Should see: "🚫 **Image not included in response**" message');
+  console.log('');
+  console.log('Expected behavior:');
+  console.log('- Screenshot file should be created');
+  console.log('- No large image sent to API (prevents conversation issues)');
+  console.log('- Clear warning message displayed');
+  
+  return testFile;
+}
+
+testLargeScreenshot().catch(console.error);
\ No newline at end of file