From 57b95acd11027929bccc32e0ccb64ea4a4ea7c9f Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Tue, 17 Sep 2024 19:29:07 -0400
Subject: [PATCH] Highlight mentions in task header and messages

---
 webview-ui/src/components/ChatRow.tsx    |  3 ++-
 webview-ui/src/components/TaskHeader.tsx | 25 ++++++++++++++----------
 webview-ui/src/index.css                 | 13 +++++++++++-
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/webview-ui/src/components/ChatRow.tsx b/webview-ui/src/components/ChatRow.tsx
index e56c68c..a87fa4b 100644
--- a/webview-ui/src/components/ChatRow.tsx
+++ b/webview-ui/src/components/ChatRow.tsx
@@ -8,6 +8,7 @@ import CodeAccordian, { formatFilePathForTruncation } from "./CodeAccordian"
 import CodeBlock, { CODE_BLOCK_BG_COLOR } from "./CodeBlock"
 import Thumbnails from "./Thumbnails"
 import { vscode } from "../utils/vscode"
+import { highlightMentions } from "./TaskHeader"
 
 interface ChatRowProps {
 	message: ClaudeMessage
@@ -425,7 +426,7 @@ const ChatRowContent = ({ message, isExpanded, onToggleExpand, lastModifiedMessa
 								whiteSpace: "pre-line",
 								wordWrap: "break-word",
 							}}>
-							<span style={{ display: "block" }}>{message.text}</span>
+							<span style={{ display: "block" }}>{highlightMentions(message.text)}</span>
 							{message.images && message.images.length > 0 && (
 								<Thumbnails images={message.images} style={{ marginTop: "8px" }} />
 							)}
diff --git a/webview-ui/src/components/TaskHeader.tsx b/webview-ui/src/components/TaskHeader.tsx
index 850adc8..bbbd496 100644
--- a/webview-ui/src/components/TaskHeader.tsx
+++ b/webview-ui/src/components/TaskHeader.tsx
@@ -148,7 +148,9 @@ 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 }}>{highlightMentions(task.text)}</span>}
+							{!isTaskExpanded && (
+								<span style={{ marginLeft: 4 }}>{highlightMentions(task.text, false)}</span>
+							)}
 						</div>
 					</div>
 					{!isTaskExpanded && isCostAvailable && (
@@ -194,7 +196,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 									wordBreak: "break-word",
 									overflowWrap: "anywhere",
 								}}>
-								{highlightMentions(task.text)}
+								{highlightMentions(task.text, false)}
 							</div>
 							{!isTextExpanded && showSeeMore && (
 								<div
@@ -337,23 +339,26 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	)
 }
 
-const highlightMentions = (text?: string) => {
-	if (!text) return []
+export const highlightMentions = (text?: string, withShadow = true) => {
+	if (!text) return text
 	const parts = text.split(mentionRegexGlobal)
-	return parts.reduce((acc, part, index) => {
+	return parts.map((part, index) => {
 		if (index % 2 === 0) {
 			// This is regular text
-			acc.push(part)
+			return part
 		} else {
 			// This is a mention
-			acc.push(
-				<span style={{ backgroundColor: "yellow" }} key={`mention-${index}`}>
+			return (
+				<span
+					key={index}
+					className={`mention-context-highlight-${
+						withShadow ? "visible-with-shadow" : "visible-without-shadow"
+					}`}>
 					@{part}
 				</span>
 			)
 		}
-		return acc
-	}, [] as (string | JSX.Element)[])
+	})
 }
 
 const ExportButton = () => (
diff --git a/webview-ui/src/index.css b/webview-ui/src/index.css
index 8ef8646..44aaa7d 100644
--- a/webview-ui/src/index.css
+++ b/webview-ui/src/index.css
@@ -146,7 +146,7 @@ vscode-dropdown::part(listbox) {
 .mention-context-highlight {
 	background-color: var(--vscode-textLink-activeForeground, #3794ff);
 	opacity: 0.3;
-	border-radius: 2px;
+	border-radius: 3px;
 	box-shadow: 0 0 0 0.5px var(--vscode-textLink-activeForeground, #3794ff);
 	color: transparent;
 	padding: 0.5px;
@@ -154,3 +154,14 @@ vscode-dropdown::part(listbox) {
 	position: relative;
 	bottom: -0.5px;
 }
+
+.mention-context-highlight-visible-with-shadow {
+	background-color: color-mix(in srgb, var(--vscode-badge-foreground) 35%, transparent);
+	border-radius: 3px;
+	box-shadow: 0 0 0 0.5px color-mix(in srgb, var(--vscode-badge-foreground) 35%, transparent);
+}
+
+.mention-context-highlight-visible-without-shadow {
+	background-color: color-mix(in srgb, var(--vscode-badge-foreground) 35%, transparent);
+	border-radius: 3px;
+}