From c846bc4455d93b2016f247db7a7495aeb384c5d3 Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Tue, 18 Feb 2025 22:14:36 -0800
Subject: [PATCH] Adjust mermaid diagram colors

---
 src/core/prompts/system.ts                    |  2 +-
 .../src/components/common/MermaidBlock.tsx    | 74 +++++++++++++++----
 2 files changed, 62 insertions(+), 14 deletions(-)

diff --git a/src/core/prompts/system.ts b/src/core/prompts/system.ts
index 1f8fcce010..8d6d13766d 100644
--- a/src/core/prompts/system.ts
+++ b/src/core/prompts/system.ts
@@ -885,7 +885,7 @@ In each user message, the environment_details will specify the current mode. The
 - When starting in PLAN MODE, depending on the user's request, you may need to do some information gathering e.g. using read_file or search_files to get more context about the task. You may also ask the user clarifying questions to get a better understanding of the task. You may return mermaid diagrams to visually display your understanding.
 - Once you've gained more context about the user's request, you should architect a detailed plan for how you will accomplish the task. Returning mermaid diagrams may be helpful here as well.
 - Then you might ask the user if they are pleased with this plan, or if they would like to make any changes. Think of this as a brainstorming session where you can discuss the task and plan the best way to accomplish it.
-- If at any point a mermaid diagram would make your plan clearer to help the user quickly see the structure, you are encouraged to include a Mermaid code block in the response.
+- If at any point a mermaid diagram would make your plan clearer to help the user quickly see the structure, you are encouraged to include a Mermaid code block in the response. (Note: if you use colors in your mermaid diagrams, be sure to use high contrast colors so the text is readable.)
 - Finally once it seems like you've reached a good plan, ask the user to switch you back to ACT MODE to implement the solution.
 
 ====
diff --git a/webview-ui/src/components/common/MermaidBlock.tsx b/webview-ui/src/components/common/MermaidBlock.tsx
index ea9d37a4f7..619188179e 100644
--- a/webview-ui/src/components/common/MermaidBlock.tsx
+++ b/webview-ui/src/components/common/MermaidBlock.tsx
@@ -5,25 +5,73 @@ import styled from "styled-components"
 import { vscode } from "../../utils/vscode"
 
 const MERMAID_THEME = {
-	background: "#1e1e1e",
-	textColor: "#ffffff",
-	mainBkg: "#2d2d2d",
-	lineColor: "#cccccc",
-	primaryColor: "#3c3c3c",
+	background: "#1e1e1e", // VS Code dark theme background
+	textColor: "#ffffff", // Main text color
+	mainBkg: "#2d2d2d", // Background for nodes
+	nodeBorder: "#888888", // Border color for nodes
+	lineColor: "#cccccc", // Lines connecting nodes
+	primaryColor: "#3c3c3c", // Primary color for highlights
+	primaryTextColor: "#ffffff", // Text in primary colored elements
+	primaryBorderColor: "#888888",
+	secondaryColor: "#2d2d2d", // Secondary color for alternate elements
+	tertiaryColor: "#454545", // Third color for special elements
+
+	// Class diagram specific
+	classText: "#ffffff",
+
+	// State diagram specific
+	labelColor: "#ffffff",
+
+	// Sequence diagram specific
+	actorLineColor: "#cccccc",
+	actorBkg: "#2d2d2d",
+	actorBorder: "#888888",
+	actorTextColor: "#ffffff",
+
+	// Flow diagram specific
+	fillType0: "#2d2d2d",
+	fillType1: "#3c3c3c",
+	fillType2: "#454545",
 }
 
 mermaid.initialize({
 	startOnLoad: false,
 	securityLevel: "loose",
 	theme: "dark",
-	// themeVariables: {
-	// 	background: MERMAID_THEME.background,
-	// 	textColor: MERMAID_THEME.textColor,
-	// 	mainBkg: MERMAID_THEME.mainBkg,
-	// 	lineColor: MERMAID_THEME.lineColor,
-	// 	fontSize: "16px",
-	// 	primaryColor: MERMAID_THEME.primaryColor,
-	// },
+	themeVariables: {
+		...MERMAID_THEME,
+		fontSize: "16px",
+		fontFamily: "var(--vscode-font-family, 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif)",
+
+		// Additional styling
+		noteTextColor: "#ffffff",
+		noteBkgColor: "#454545",
+		noteBorderColor: "#888888",
+
+		// Improve contrast for special elements
+		critBorderColor: "#ff9580",
+		critBkgColor: "#803d36",
+
+		// Task diagram specific
+		taskTextColor: "#ffffff",
+		taskTextOutsideColor: "#ffffff",
+		taskTextLightColor: "#ffffff",
+
+		// Numbers/sections
+		sectionBkgColor: "#2d2d2d",
+		sectionBkgColor2: "#3c3c3c",
+
+		// Alt sections in sequence diagrams
+		altBackground: "#2d2d2d",
+
+		// Links
+		linkColor: "#6cb6ff",
+
+		// Borders and lines
+		compositeBackground: "#2d2d2d",
+		compositeBorder: "#888888",
+		titleColor: "#ffffff",
+	},
 })
 
 interface MermaidBlockProps {