From 706ee8a0fee25fb05c3697f679dd343c44f413ba Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Tue, 10 Sep 2024 17:18:21 -0400
Subject: [PATCH] Add ability to open images

---
 src/providers/ClaudeDevProvider.ts           |  4 ++++
 src/shared/WebviewMessage.ts                 |  1 +
 src/utils/open-image.ts                      | 20 ++++++++++++++++++++
 webview-ui/src/components/HistoryPreview.tsx |  2 +-
 webview-ui/src/components/TaskHeader.tsx     |  2 +-
 webview-ui/src/components/Thumbnails.tsx     |  7 +++++++
 6 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 src/utils/open-image.ts

diff --git a/src/providers/ClaudeDevProvider.ts b/src/providers/ClaudeDevProvider.ts
index ba0f037..8d1d513 100644
--- a/src/providers/ClaudeDevProvider.ts
+++ b/src/providers/ClaudeDevProvider.ts
@@ -10,6 +10,7 @@ import fs from "fs/promises"
 import { HistoryItem } from "../shared/HistoryItem"
 import axios from "axios"
 import { getTheme } from "../utils/getTheme"
+import { openImage } from "../utils/open-image"
 
 /*
 https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default/weather-webview/src/providers/WeatherViewProvider.ts
@@ -398,6 +399,9 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 						const models = await this.getOllamaModels(message.text)
 						this.postMessageToWebview({ type: "ollamaModels", models })
 						break
+					case "openImage":
+						openImage(message.text!)
+						break
 					// Add more switch case statements here as more webview message commands
 					// are created within the webview context (i.e. inside media/main.js)
 				}
diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts
index 804e0aa..c3a35f0 100644
--- a/src/shared/WebviewMessage.ts
+++ b/src/shared/WebviewMessage.ts
@@ -17,6 +17,7 @@ export interface WebviewMessage {
 		| "exportTaskWithId"
 		| "resetState"
 		| "requestOllamaModels"
+		| "openImage"
 	text?: string
 	askResponse?: ClaudeAskResponse
 	apiConfiguration?: ApiConfiguration
diff --git a/src/utils/open-image.ts b/src/utils/open-image.ts
new file mode 100644
index 0000000..42e09ee
--- /dev/null
+++ b/src/utils/open-image.ts
@@ -0,0 +1,20 @@
+import * as path from "path"
+import * as os from "os"
+import * as vscode from "vscode"
+
+export async function openImage(dataUri: string) {
+	const matches = dataUri.match(/^data:image\/([a-zA-Z]+);base64,(.+)$/)
+	if (!matches) {
+		vscode.window.showErrorMessage("Invalid data URI format")
+		return
+	}
+	const [, format, base64Data] = matches
+	const imageBuffer = Buffer.from(base64Data, "base64")
+	const tempFilePath = path.join(os.tmpdir(), `temp_image_${Date.now()}.${format}`)
+	try {
+		await vscode.workspace.fs.writeFile(vscode.Uri.file(tempFilePath), imageBuffer)
+		await vscode.commands.executeCommand("vscode.open", vscode.Uri.file(tempFilePath))
+	} catch (error) {
+		vscode.window.showErrorMessage(`Error opening image: ${error}`)
+	}
+}
diff --git a/webview-ui/src/components/HistoryPreview.tsx b/webview-ui/src/components/HistoryPreview.tsx
index a9aa2c7..b13d431 100644
--- a/webview-ui/src/components/HistoryPreview.tsx
+++ b/webview-ui/src/components/HistoryPreview.tsx
@@ -33,7 +33,7 @@ const HistoryPreview = ({ showHistoryView }: HistoryPreviewProps) => {
 			<style>
 				{`
 					.history-preview-item {
-						background-color: color-mix(in srgb, var(--vscode-toolbar-hoverBackground) 50%, transparent);
+						background-color: color-mix(in srgb, var(--vscode-toolbar-hoverBackground) 65%, transparent);
 						border-radius: 4px;
 						position: relative;
 						overflow: hidden;
diff --git a/webview-ui/src/components/TaskHeader.tsx b/webview-ui/src/components/TaskHeader.tsx
index 2bcd82e..c7b66be 100644
--- a/webview-ui/src/components/TaskHeader.tsx
+++ b/webview-ui/src/components/TaskHeader.tsx
@@ -156,7 +156,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 							${totalCost?.toFixed(4)}
 						</div>
 					)}
-					<VSCodeButton appearance="icon" onClick={onClose} style={{ marginLeft: 10, flexShrink: 0 }}>
+					<VSCodeButton appearance="icon" onClick={onClose} style={{ marginLeft: 6, flexShrink: 0 }}>
 						<span className="codicon codicon-close"></span>
 					</VSCodeButton>
 				</div>
diff --git a/webview-ui/src/components/Thumbnails.tsx b/webview-ui/src/components/Thumbnails.tsx
index b548600..31ca437 100644
--- a/webview-ui/src/components/Thumbnails.tsx
+++ b/webview-ui/src/components/Thumbnails.tsx
@@ -1,5 +1,6 @@
 import React, { useState, useRef, useLayoutEffect, memo } from "react"
 import { useWindowSize } from "react-use"
+import { vscode } from "../utils/vscode"
 
 interface ThumbnailsProps {
 	images: string[]
@@ -31,6 +32,10 @@ const Thumbnails = ({ images, style, setImages, onHeightChange }: ThumbnailsProp
 
 	const isDeletable = setImages !== undefined
 
+	const handleImageClick = (image: string) => {
+		vscode.postMessage({ type: "openImage", text: image })
+	}
+
 	return (
 		<div
 			ref={containerRef}
@@ -55,7 +60,9 @@ const Thumbnails = ({ images, style, setImages, onHeightChange }: ThumbnailsProp
 							height: 34,
 							objectFit: "cover",
 							borderRadius: 4,
+							cursor: "pointer",
 						}}
+						onClick={() => handleImageClick(image)}
 					/>
 					{isDeletable && hoveredIndex === index && (
 						<div