From 30ed938dedc49cff4d82100a6623894d4425c5a5 Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Tue, 17 Sep 2024 18:51:33 -0400
Subject: [PATCH] Refactor mention regex; highlight task text mentions

---
 webview-ui/src/components/ChatTextArea.tsx | 16 +++++++++++----
 webview-ui/src/components/TaskHeader.tsx   | 24 ++++++++++++++++++++--
 webview-ui/src/utils/mention-context.ts    | 23 ++++++++-------------
 3 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/webview-ui/src/components/ChatTextArea.tsx b/webview-ui/src/components/ChatTextArea.tsx
index 59bd63f..fb9e5f7 100644
--- a/webview-ui/src/components/ChatTextArea.tsx
+++ b/webview-ui/src/components/ChatTextArea.tsx
@@ -1,7 +1,14 @@
 import React, { forwardRef, useCallback, useEffect, useLayoutEffect, useRef, useState } from "react"
 import DynamicTextArea from "react-textarea-autosize"
 import { useExtensionState } from "../context/ExtensionStateContext"
-import { getContextMenuOptions, insertMention, removeMention, shouldShowContextMenu } from "../utils/mention-context"
+import {
+	getContextMenuOptions,
+	insertMention,
+	mentionRegex,
+	mentionRegexGlobal,
+	removeMention,
+	shouldShowContextMenu,
+} from "../utils/mention-context"
 import { MAX_IMAGES_PER_MESSAGE } from "./ChatView"
 import ContextMenu from "./ContextMenu"
 import Thumbnails from "./Thumbnails"
@@ -184,11 +191,13 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 						charBeforeCursor === " " || charBeforeCursor === "\n" || charBeforeCursor === "\r\n"
 					const charAfterIsWhitespace =
 						charAfterCursor === " " || charAfterCursor === "\n" || charAfterCursor === "\r\n"
+					// checks if char before cusor is whitespace after a mention
 					if (
 						charBeforeIsWhitespace &&
-						inputValue.slice(0, cursorPosition - 1).match(/@((?:\/|\w+:\/\/)[^\s]+|problems)$/)
+						inputValue.slice(0, cursorPosition - 1).match(new RegExp(mentionRegex.source + "$")) // "$" is added to ensure the match occurs at the end of the string
 					) {
 						const newCursorPosition = cursorPosition - 1
+						// if mention is followed by another word, then instead of deleting the space separating them we just move the cursor to the end of the mention
 						if (!charAfterIsWhitespace) {
 							event.preventDefault()
 							textAreaRef.current?.setSelectionRange(newCursorPosition, newCursorPosition)
@@ -349,12 +358,11 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
 			if (!textAreaRef.current || !highlightLayerRef.current) return
 
 			const text = textAreaRef.current.value
-			const mentionRegex = /@((?:\/|\w+:\/\/)[^\s]+|problems\b)/g
 
 			highlightLayerRef.current.innerHTML = text
 				.replace(/\n$/, "\n\n")
 				.replace(/[<>&]/g, (c) => ({ "<": "&lt;", ">": "&gt;", "&": "&amp;" }[c] || c))
-				.replace(mentionRegex, '<mark class="mention-context-highlight">$&</mark>')
+				.replace(mentionRegexGlobal, '<mark class="mention-context-highlight">$&</mark>')
 
 			highlightLayerRef.current.scrollTop = textAreaRef.current.scrollTop
 			highlightLayerRef.current.scrollLeft = textAreaRef.current.scrollLeft
diff --git a/webview-ui/src/components/TaskHeader.tsx b/webview-ui/src/components/TaskHeader.tsx
index d776251..850adc8 100644
--- a/webview-ui/src/components/TaskHeader.tsx
+++ b/webview-ui/src/components/TaskHeader.tsx
@@ -3,6 +3,7 @@ import React, { memo, useEffect, useMemo, useRef, useState } from "react"
 import { useWindowSize } from "react-use"
 import { ClaudeMessage } from "../../../src/shared/ExtensionMessage"
 import { useExtensionState } from "../context/ExtensionStateContext"
+import { mentionRegexGlobal } from "../utils/mention-context"
 import { vscode } from "../utils/vscode"
 import Thumbnails from "./Thumbnails"
 
@@ -147,7 +148,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 								minWidth: 0, // This allows the div to shrink below its content size
 							}}>
 							<span style={{ fontWeight: "bold" }}>Task{!isTaskExpanded && ":"}</span>
-							{!isTaskExpanded && <span style={{ marginLeft: 4 }}>{task.text}</span>}
+							{!isTaskExpanded && <span style={{ marginLeft: 4 }}>{highlightMentions(task.text)}</span>}
 						</div>
 					</div>
 					{!isTaskExpanded && isCostAvailable && (
@@ -193,7 +194,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 									wordBreak: "break-word",
 									overflowWrap: "anywhere",
 								}}>
-								{task.text}
+								{highlightMentions(task.text)}
 							</div>
 							{!isTextExpanded && showSeeMore && (
 								<div
@@ -336,6 +337,25 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	)
 }
 
+const highlightMentions = (text?: string) => {
+	if (!text) return []
+	const parts = text.split(mentionRegexGlobal)
+	return parts.reduce((acc, part, index) => {
+		if (index % 2 === 0) {
+			// This is regular text
+			acc.push(part)
+		} else {
+			// This is a mention
+			acc.push(
+				<span style={{ backgroundColor: "yellow" }} key={`mention-${index}`}>
+					@{part}
+				</span>
+			)
+		}
+		return acc
+	}, [] as (string | JSX.Element)[])
+}
+
 const ExportButton = () => (
 	<VSCodeButton
 		appearance="icon"
diff --git a/webview-ui/src/utils/mention-context.ts b/webview-ui/src/utils/mention-context.ts
index ee204d5..2809cea 100644
--- a/webview-ui/src/utils/mention-context.ts
+++ b/webview-ui/src/utils/mention-context.ts
@@ -1,16 +1,12 @@
-// export const mockPaths = [
-// 	{ type: "problems", path: "Problems" },
-// 	{ type: "file", path: "/src/components/Header.tsx" },
-// 	{ type: "file", path: "/src/components/Footer.tsx" },
-// 	{ type: "file", path: "/src/utils/helpers.ts" },
-// 	{ type: "folder", path: "/src/components" },
-// 	{ type: "folder", path: "/src/utils" },
-// 	{ type: "folder", path: "/public/images" },
-// 	{ type: "file", path: "/public/index.html" },
-// 	{ type: "file", path: "/package.json" },
-// 	{ type: "folder", path: "/node_modules" },
-// 	{ type: "file", path: "/README.md" },
-// ]
+/*
+Mention regex
+- File and folder paths (starting with '/')
+- URLs (containing '://')
+- The 'problems' keyword
+- Word boundary after 'problems' to avoid partial matches
+*/
+export const mentionRegex = /@((?:\/|\w+:\/\/)[^\s]+|problems\b)/
+export const mentionRegexGlobal = new RegExp(mentionRegex.source, "g")
 
 export function insertMention(text: string, position: number, value: string): string {
 	const beforeCursor = text.slice(0, position)
@@ -30,7 +26,6 @@ export function insertMention(text: string, position: number, value: string): st
 }
 
 export function removeMention(text: string, position: number): { newText: string; newPosition: number } {
-	const mentionRegex = /@((?:\/|\w+:\/\/)[^\s]+|problems\b)/
 	const beforeCursor = text.slice(0, position)
 	const afterCursor = text.slice(position)