From 46ceafa016c77b1ccce4d0507aa7ff2cd4cc0086 Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Sat, 21 Sep 2024 16:24:59 -0400
Subject: [PATCH] Close browser even if screenshot fails; handle case where
 webp screenshot fails

---
 src/ClaudeDev.ts               | 12 ++++++++++--
 src/utils/UrlContentFetcher.ts | 34 +++++++++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/ClaudeDev.ts b/src/ClaudeDev.ts
index 9182249..cf2f238 100644
--- a/src/ClaudeDev.ts
+++ b/src/ClaudeDev.ts
@@ -1471,8 +1471,16 @@ export class ClaudeDev {
 
 			await this.say("inspect_site_result", "") // no result, starts the loading spinner waiting for result
 			await this.urlContentFetcher.launchBrowser()
-			const { screenshot, logs } = await this.urlContentFetcher.urlToScreenshotAndLogs(url)
-			await this.urlContentFetcher.closeBrowser()
+			let result: {
+				screenshot: string
+				logs: string
+			}
+			try {
+				result = await this.urlContentFetcher.urlToScreenshotAndLogs(url)
+			} finally {
+				await this.urlContentFetcher.closeBrowser()
+			}
+			const { screenshot, logs } = result
 			await this.say("inspect_site_result", logs, [screenshot])
 
 			return [
diff --git a/src/utils/UrlContentFetcher.ts b/src/utils/UrlContentFetcher.ts
index a3408e9..db02940 100644
--- a/src/utils/UrlContentFetcher.ts
+++ b/src/utils/UrlContentFetcher.ts
@@ -1,7 +1,7 @@
 import * as vscode from "vscode"
 import * as fs from "fs/promises"
 import * as path from "path"
-import { Browser, Page, TimeoutError, launch } from "puppeteer-core"
+import { Browser, Page, ScreenshotOptions, TimeoutError, launch } from "puppeteer-core"
 import * as cheerio from "cheerio"
 import TurndownService from "turndown"
 // @ts-ignore
@@ -130,14 +130,38 @@ export class UrlContentFetcher {
 			interval: 100,
 		}).catch(() => {})
 
-		const screenshotBase64 = await this.page.screenshot({
+		// image cannot exceed 8_000 pixels
+		const pageHeight = await this.page.evaluate(() => document.documentElement.scrollHeight)
+		let options: ScreenshotOptions = {
 			fullPage: true,
-			type: "webp",
 			encoding: "base64",
 			// quality: 80,
-		})
+			clip: {
+				x: 0,
+				y: 0,
+				width: await this.page.evaluate(() => document.documentElement.clientWidth),
+				height: Math.min(pageHeight, 8_000),
+			},
+		}
 
-		const screenshot = `data:image/webp;base64,${screenshotBase64}`
+		let screenshotBase64 = await this.page.screenshot({
+			...options,
+			type: "webp",
+		})
+		let screenshot = `data:image/webp;base64,${screenshotBase64}`
+
+		if (!screenshotBase64) {
+			console.log("webp screenshot failed, trying png")
+			screenshotBase64 = await this.page.screenshot({
+				...options,
+				type: "png",
+			})
+			screenshot = `data:image/png;base64,${screenshotBase64}`
+		}
+
+		if (!screenshotBase64) {
+			throw new Error("Failed to take screenshot.")
+		}
 
 		this.page.removeAllListeners()