refactor: memoize renderSegment with useCallback

webview-ui/src/components/mcp/chat-display/McpResponseDisplay.tsx

@@ -6,7 +6,7 @@ import ImagePreview from "./ImagePreview"
 import styled from "styled-components"
 import { CODE_BLOCK_BG_COLOR } from "@/components/common/CodeBlock"
 import ChatErrorBoundary from "@/components/chat/ChatErrorBoundary"
-import { isUrl, isLocalhostUrl, formatUrlForOpening, checkIfImageUrl } from "./utils/mcpRichUtil"
+import { UrlMatch, processResponseUrls, DisplaySegment, buildDisplaySegments } from "./utils/mcpRichUtil"
 
 // Maximum number of URLs to process in total, per response
 export const MAX_URLS = 50
@@ -107,23 +107,12 @@ interface McpResponseDisplayProps {
 	responseText: string
 }
 
-// Represents a URL found in the text with its position and metadata
-interface UrlMatch {
-	url: string // The actual URL
-	fullMatch: string // The full matched text
-	index: number // Position in the text
-	isImage: boolean // Whether this URL is an image
-	isProcessed: boolean // Whether we've already processed this URL (to avoid duplicates)
-}
-
 const McpResponseDisplay: React.FC<McpResponseDisplayProps> = ({ responseText }) => {
 	const { mcpResponsesCollapsed, mcpRichDisplayEnabled } = useExtensionState() // Get setting from context
 	const [isExpanded, setIsExpanded] = useState(!mcpResponsesCollapsed) // Initialize with context setting
 	const [isLoading, setIsLoading] = useState(false) // Initial loading state for rich content
-	const [displayMode, setDisplayMode] = useState<"rich" | "plain">(() => {
-		// Initialize directly from the global setting.
-		return mcpRichDisplayEnabled ? "rich" : "plain"
-	})
+	const [displayMode, setDisplayMode] = useState<"rich" | "plain">(mcpRichDisplayEnabled ? "rich" : "plain")
+
 	const [urlMatches, setUrlMatches] = useState<UrlMatch[]>([])
 	const [error, setError] = useState<string | null>(null)
 
@@ -151,126 +140,77 @@ const McpResponseDisplay: React.FC<McpResponseDisplayProps> = ({ responseText })
 			return
 		}
 
-		// Use a direct boolean for cancellation that's scoped to this effect run
-		let processingCanceled = false
-		const processResponse = async () => {
-			console.log("Processing MCP response for URL extraction")
-			setIsLoading(true)
-			setError(null)
-			try {
-				const text = responseText || ""
-				const matches: UrlMatch[] = []
-				const urlRegex = /(?:https?:\/\/|data:image)[^\s<>"']+/g
-				let urlMatch: RegExpExecArray | null
-				let urlCount = 0
+		console.log("Processing MCP response for URL extraction")
+		setIsLoading(true)
+		setError(null)
 
-				// First pass: Extract all URLs and immediately make them available for rendering
-				while ((urlMatch = urlRegex.exec(text)) !== null && urlCount < MAX_URLS) {
-					// Get the original URL from the match - never modify the original URL text
-					const url = urlMatch[0]
-
-					// Skip invalid URLs
-					if (!isUrl(url)) {
-						console.log("Skipping invalid URL:", url)
-						continue
-					}
-
-					// Skip localhost URLs to prevent security issues
-					if (isLocalhostUrl(url)) {
-						console.log("Skipping localhost URL:", url)
-						continue
-					}
-
-					matches.push({
-						url,
-						fullMatch: url,
-						index: urlMatch.index,
-						isImage: false, // Will check later
-						isProcessed: false,
-					})
-
-					urlCount++
-				}
-
-				console.log(`Found ${matches.length} URLs in text, will check if they are images`)
-
-				// Set matches immediately so UI can start rendering with loading states
-				setUrlMatches(matches.sort((a, b) => a.index - b.index))
-
-				// Mark loading as complete to show content immediately
+		// Use the orchestrator function from mcpRichUtil
+		const cleanup = processResponseUrls(
+			responseText || "",
+			MAX_URLS,
+			(matches) => {
+				setUrlMatches(matches)
 				setIsLoading(false)
-
-				// Process image checks in the background - one at a time to avoid network flooding
-				const processImageChecks = async () => {
-					console.log(`Starting sequential URL processing for ${matches.length} URLs`)
-
-					for (let i = 0; i < matches.length; i++) {
-						// Skip already processed URLs (from extension check)
-						if (matches[i].isProcessed) continue
-
-						// Check if processing has been canceled (switched to plain mode)
-						if (processingCanceled) {
-							console.log("URL processing canceled - display mode changed to plain")
-							return
-						}
-
-						const match = matches[i]
-						console.log(`Processing URL ${i + 1} of ${matches.length}: ${match.url}`)
-
-						try {
-							// Process each URL individually
-							const isImage = await checkIfImageUrl(match.url)
-
-							// Skip if processing has been canceled
-							if (processingCanceled) return
-
-							// Update the match in place
-							match.isImage = isImage
-							match.isProcessed = true
-
-							// Update state after each URL to show progress
-							// Create a new array to ensure React detects the state change
-							setUrlMatches([...matches])
-						} catch (err) {
-							console.log(`URL check error: ${match.url}`, err)
-							match.isProcessed = true
-
-							// Update state even on error
-							if (!processingCanceled) {
-								setUrlMatches([...matches])
-							}
-						}
-
-						// Delay between URL processing to avoid overwhelming the network
-						if (!processingCanceled && i < matches.length - 1) {
-							await new Promise((resolve) => setTimeout(resolve, 100))
-						}
-					}
-
-					console.log(`URL processing complete. Found ${matches.filter((m) => m.isImage).length} image URLs`)
-				}
-
-				// Start the background processing
-				processImageChecks()
-			} catch (error) {
-				setError("Failed to process response content. Switch to plain text mode to view safely.")
+			},
+			(updatedMatches) => {
+				setUrlMatches(updatedMatches)
+			},
+			(errorMessage) => {
+				setError(errorMessage)
 				setIsLoading(false)
-			}
-		}
+			},
+		)
 
-		processResponse()
-
-		// Cleanup function to cancel processing if component unmounts or dependencies change
-		return () => {
-			processingCanceled = true
-			console.log("Cleaning up URL processing")
-		}
+		return cleanup
 	}, [responseText, displayMode, isExpanded])
 
+	// Helper function to render a display segment
+	const renderSegment = useCallback((segment: DisplaySegment): JSX.Element => {
+		switch (segment.type) {
+			case "text":
+			case "url":
+				return <UrlText key={segment.key}>{segment.content}</UrlText>
+
+			case "image":
+				return (
+					<div key={segment.key}>
+						<ImagePreview url={segment.url!} />
+					</div>
+				)
+
+			case "link":
+				return (
+					<div key={segment.key} style={{ margin: "10px 0" }}>
+						<LinkPreview url={segment.url!} />
+					</div>
+				)
+
+			case "error":
+				return (
+					<div
+						key={segment.key}
+						style={{
+							margin: "10px 0",
+							padding: "8px",
+							color: "var(--vscode-errorForeground)",
+							border: "1px solid var(--vscode-editorError-foreground)",
+							borderRadius: "4px",
+							height: "128px",
+							overflow: "auto",
+						}}>
+						{segment.content}
+					</div>
+				)
+
+			default:
+				return <React.Fragment key={segment.key} />
+		}
+	}, [])
+
 	// Function to render content based on display mode
 	const renderContent = () => {
 		if (!isExpanded) {
-			return null // Don't render content if not expanded
+			return null
 		}
 
 		if (isLoading && displayMode === "rich") {
@@ -281,12 +221,10 @@ const McpResponseDisplay: React.FC<McpResponseDisplayProps> = ({ responseText })
 			)
 		}
 
-		// For plain text mode, just show the text
 		if (displayMode === "plain") {
 			return <UrlText>{responseText}</UrlText>
 		}
 
-		// Show error message if there was an error
 		if (error) {
 			return (
 				<>
@@ -296,97 +234,9 @@ const McpResponseDisplay: React.FC<McpResponseDisplayProps> = ({ responseText })
 			)
 		}
 
-		// For rich display mode, show the text with embedded content
 		if (displayMode === "rich") {
-			// We already know displayMode is "rich" if we get here
-			// Create an array of text segments and embedded content
-			const segments: JSX.Element[] = []
-			let lastIndex = 0
-			let segmentIndex = 0
-
-			// Track embed count for logging
-			let embedCount = 0
-
-			// Add the text before the first URL
-			if (urlMatches.length === 0) {
-				segments.push(<UrlText key={`segment-${segmentIndex}`}>{responseText}</UrlText>)
-			} else {
-				for (let i = 0; i < urlMatches.length; i++) {
-					const match = urlMatches[i]
-					const { url, fullMatch, index } = match
-
-					// Add text segment before this URL
-					if (index > lastIndex) {
-						segments.push(
-							<UrlText key={`segment-${segmentIndex++}`}>{responseText.substring(lastIndex, index)}</UrlText>,
-						)
-					}
-
-					// Add the URL text itself
-					segments.push(<UrlText key={`url-${segmentIndex++}`}>{fullMatch}</UrlText>)
-
-					// Calculate the end position of this URL in the text
-					const urlEndIndex = index + fullMatch.length
-
-					// Add embedded content after the URL
-					// For images, use the ImagePreview component
-					if (match.isImage) {
-						segments.push(
-							<div key={`embed-image-${url}-${segmentIndex++}`}>
-								{/* Use formatUrlForOpening for network calls but preserve original URL in display */}
-								<ImagePreview url={formatUrlForOpening(url)} />
-							</div>,
-						)
-						embedCount++
-						// console.log(`Added image embed for ${url}, embed count: ${embedCount}`);
-					} else if (match.isProcessed) {
-						// For non-image URLs or URLs we haven't processed yet, show link preview
-						try {
-							// Skip localhost URLs
-							if (!isLocalhostUrl(url)) {
-								// Use a unique key that includes the URL to ensure each preview is isolated
-								segments.push(
-									<div key={`embed-${url}-${segmentIndex++}`} style={{ margin: "10px 0" }}>
-										{/* Already using formatUrlForOpening for link previews */}
-										<LinkPreview url={formatUrlForOpening(url)} />
-									</div>,
-								)
-
-								embedCount++
-								// console.log(`Added link preview for ${url}, embed count: ${embedCount}`);
-							}
-						} catch (e) {
-							console.log("Link preview could not be created")
-							// Show error message for failed link preview
-							segments.push(
-								<div
-									key={`embed-error-${segmentIndex++}`}
-									style={{
-										margin: "10px 0",
-										padding: "8px",
-										color: "var(--vscode-errorForeground)",
-										border: "1px solid var(--vscode-editorError-foreground)",
-										borderRadius: "4px",
-										height: "128px", // Fixed height
-										overflow: "auto", // Allow scrolling if content overflows
-									}}>
-									Failed to create preview for: {url}
-								</div>,
-							)
-						}
-					}
-
-					// Update lastIndex for next segment
-					lastIndex = urlEndIndex
-				}
-
-				// Add any remaining text after the last URL
-				if (lastIndex < responseText.length) {
-					segments.push(<UrlText key={`segment-${segmentIndex++}`}>{responseText.substring(lastIndex)}</UrlText>)
-				}
-			}
-
-			return <>{segments}</>
+			const segments = buildDisplaySegments(responseText, urlMatches)
+			return <>{segments.map(renderSegment)}</>
 		}
 
 		return null

webview-ui/src/components/mcp/chat-display/utils/mcpRichUtil.ts

@@ -1,6 +1,23 @@
 import { WebServiceClient } from "@/services/grpc-client"
 import { StringRequest } from "@shared/proto/common"
 
+// Represents a URL found in the text with its position and metadata
+export interface UrlMatch {
+	url: string // The actual URL
+	fullMatch: string // The full matched text
+	index: number // Position in the text
+	isImage: boolean // Whether this URL is an image
+	isProcessed: boolean // Whether we've already processed this URL (to avoid duplicates)
+}
+
+// Display segment interface
+export interface DisplaySegment {
+	type: "text" | "url" | "image" | "link" | "error"
+	content: string
+	url?: string
+	key: string // Pre-computed key for React
+}
+
 // Safely create a URL object with error handling and ensure HTTPS
 export const safeCreateUrl = (url: string): URL | null => {
 	try {
@@ -168,3 +185,224 @@ export const checkIfImageUrl = async (url: string): Promise<boolean> => {
 	console.log(`URL protocol not supported for image check: ${url}`)
 	return false
 }
+
+/**
+ * Extracts all valid URLs from the given text
+ * @param text - The text to search for URLs
+ * @param maxUrls - Maximum number of URLs to extract (default: 50)
+ * @returns Array of URL matches sorted by position in text
+ */
+export const extractUrlsFromText = (text: string, maxUrls: number = 50): UrlMatch[] => {
+	const matches: UrlMatch[] = []
+	const urlRegex = /(?:https?:\/\/|data:image)[^\s<>"']+/g
+	let urlMatch: RegExpExecArray | null
+	let urlCount = 0
+
+	while ((urlMatch = urlRegex.exec(text)) !== null && urlCount < maxUrls) {
+		const url = urlMatch[0]
+
+		// Skip invalid URLs
+		if (!isUrl(url)) {
+			console.log("Skipping invalid URL:", url)
+			continue
+		}
+
+		// Skip localhost URLs to prevent security issues
+		if (isLocalhostUrl(url)) {
+			console.log("Skipping localhost URL:", url)
+			continue
+		}
+
+		matches.push({
+			url,
+			fullMatch: url,
+			index: urlMatch.index,
+			isImage: false, // Will be determined later
+			isProcessed: false,
+		})
+
+		urlCount++
+	}
+
+	console.log(`Found ${matches.length} URLs in text`)
+	return matches.sort((a, b) => a.index - b.index)
+}
+
+/**
+ * Processes URLs to determine their types (e.g., image vs link)
+ * Processes URLs sequentially to avoid network flooding
+ * @param matches - Array of URL matches to process
+ * @param onProgress - Callback for progress updates with updated matches
+ * @param cancellationToken - Object to check if processing should be cancelled
+ * @returns Promise that resolves when processing is complete
+ */
+export const processUrlTypes = async (
+	matches: UrlMatch[],
+	onProgress: (updatedMatches: UrlMatch[]) => void,
+	cancellationToken: { cancelled: boolean },
+): Promise<void> => {
+	console.log(`Starting sequential URL processing for ${matches.length} URLs`)
+
+	for (let i = 0; i < matches.length; i++) {
+		// Skip already processed URLs
+		if (matches[i].isProcessed) continue
+
+		// Check if processing has been canceled
+		if (cancellationToken.cancelled) {
+			console.log("URL processing canceled")
+			return
+		}
+
+		const match = matches[i]
+		console.log(`Processing URL ${i + 1} of ${matches.length}: ${match.url}`)
+
+		try {
+			// Check if URL is an image
+			const isImage = await checkIfImageUrl(match.url)
+
+			// Skip if processing has been canceled
+			if (cancellationToken.cancelled) return
+
+			// Update the match
+			match.isImage = isImage
+			match.isProcessed = true
+
+			// Notify progress with a new array to ensure React detects changes
+			onProgress([...matches])
+		} catch (err) {
+			console.log(`URL check error: ${match.url}`, err)
+			match.isProcessed = true
+
+			// Update state even on error
+			if (!cancellationToken.cancelled) {
+				onProgress([...matches])
+			}
+		}
+
+		// Delay between URL processing to avoid overwhelming the network
+		if (!cancellationToken.cancelled && i < matches.length - 1) {
+			await new Promise((resolve) => setTimeout(resolve, 100))
+		}
+	}
+
+	console.log(`URL processing complete. Found ${matches.filter((m) => m.isImage).length} image URLs`)
+}
+
+/**
+ * Orchestrates the URL extraction and processing pipeline
+ * @param text - The response text to process
+ * @param maxUrls - Maximum number of URLs to process
+ * @param onMatchesFound - Callback when initial URLs are extracted
+ * @param onMatchesUpdated - Callback when URL types are determined
+ * @param onError - Error handler callback
+ * @returns Cleanup function to cancel processing
+ */
+export const processResponseUrls = (
+	text: string,
+	maxUrls: number,
+	onMatchesFound: (matches: UrlMatch[]) => void,
+	onMatchesUpdated: (matches: UrlMatch[]) => void,
+	onError: (error: string) => void,
+): (() => void) => {
+	const cancellationToken = { cancelled: false }
+
+	const process = async () => {
+		try {
+			// Extract URLs from text
+			const matches = extractUrlsFromText(text, maxUrls)
+
+			// Immediately notify about found matches
+			onMatchesFound(matches)
+
+			// Process URLs in the background
+			await processUrlTypes(matches, onMatchesUpdated, cancellationToken)
+		} catch (error) {
+			onError("Failed to process response content. Switch to plain text mode to view safely.")
+		}
+	}
+
+	// Start processing
+	process()
+
+	// Return cleanup function
+	return () => {
+		cancellationToken.cancelled = true
+		console.log("Cleaning up URL processing")
+	}
+}
+
+/**
+ * Builds an array of display segments from response text and URL matches
+ * @param responseText - The full response text
+ * @param urlMatches - Array of URL matches with their positions and types
+ * @returns Array of display segments describing how to render the content
+ */
+export const buildDisplaySegments = (responseText: string, urlMatches: UrlMatch[]): DisplaySegment[] => {
+	const segments: DisplaySegment[] = []
+	let lastIndex = 0
+	let segmentIndex = 0
+
+	// Handle case with no URLs
+	if (urlMatches.length === 0) {
+		return [
+			{
+				type: "text",
+				content: responseText,
+				key: "segment-0",
+			},
+		]
+	}
+
+	// Process each URL match
+	for (let i = 0; i < urlMatches.length; i++) {
+		const match = urlMatches[i]
+		const { url, fullMatch, index } = match
+
+		// Add text segment before this URL
+		if (index > lastIndex) {
+			segments.push({
+				type: "text",
+				content: responseText.substring(lastIndex, index),
+				key: `segment-${segmentIndex++}`,
+			})
+		}
+
+		// Add the URL text itself
+		segments.push({
+			type: "url",
+			content: fullMatch,
+			key: `url-${segmentIndex++}`,
+		})
+
+		// Add embedded content after the URL
+		if (match.isImage) {
+			segments.push({
+				type: "image",
+				content: url,
+				url: formatUrlForOpening(url),
+				key: `embed-image-${url}-${segmentIndex++}`,
+			})
+		} else if (match.isProcessed && !isLocalhostUrl(url)) {
+			segments.push({
+				type: "link",
+				content: url,
+				url: formatUrlForOpening(url),
+				key: `embed-${url}-${segmentIndex++}`,
+			})
+		}
+
+		// Update lastIndex for next segment
+		lastIndex = index + fullMatch.length
+	}
+
+	// Add any remaining text after the last URL
+	if (lastIndex < responseText.length) {
+		segments.push({
+			type: "text",
+			content: responseText.substring(lastIndex),
+			key: `segment-${segmentIndex++}`,
+		})
+	}
+
+	return segments
+}

webview-ui/src/context/ExtensionStateContext.tsx

@@ -61,7 +61,6 @@ interface ExtensionStateContextType extends ExtensionState {
 	setPlanActSeparateModelsSetting: (value: boolean) => void
 	setEnableCheckpointsSetting: (value: boolean) => void
 	setMcpMarketplaceEnabled: (value: boolean) => void
-	setMcpRichDisplayEnabled: (value: boolean) => void
 	setMcpResponsesCollapsed: (value: boolean) => void
 	setShellIntegrationTimeout: (value: number) => void
 	setTerminalReuseEnabled: (value: boolean) => void
@@ -701,11 +700,6 @@ export const ExtensionStateContextProvider: React.FC<{
 				...prevState,
 				mcpMarketplaceEnabled: value,
 			})),
-		setMcpRichDisplayEnabled: (value) =>
-			setState((prevState) => ({
-				...prevState,
-				mcpRichDisplayEnabled: value,
-			})),
 		setMcpResponsesCollapsed: (value) => {
 			setState((prevState) => ({
 				...prevState,