Highlight mentions in task header and messages

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" }} />
 							)}

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 = () => (