update, add copy in svg

src/components/chat-view/Markdown/MarkdownWithIcon.tsx

@@ -1,13 +1,13 @@
 import * as Tooltip from '@radix-ui/react-tooltip';
 import { Check, CircleCheckBig, CircleHelp, CopyIcon, FilePlus2 } from 'lucide-react';
 import { ReactNode, useState } from 'react';
-import ReactMarkdown from 'react-markdown';
-import rehypeRaw from 'rehype-raw';
 import { useApp } from 'src/contexts/AppContext';
 
 import { t } from '../../../lang/helpers'
 
-function CopyButton({ message }: { message: string }) {
+import RawMarkdownBlock from './RawMarkdownBlock'
+
+export function CopyButton({ message }: { message: string }) {
 	const [copied, setCopied] = useState(false)
 
 	const handleCopy = async () => {
@@ -43,7 +43,7 @@ function CopyButton({ message }: { message: string }) {
 	)
 }
 
-function CreateNewFileButton({ message }: { message: string }) {
+export function CreateNewFileButton({ message }: { message: string }) {
 	const app = useApp()
 	const [created, setCreated] = useState(false)
 
@@ -138,12 +138,10 @@ const MarkdownWithIcons = ({
 		<>
 			<div className={`${className}`}>
 				<span>{iconName && renderIcon()} {renderTitle()}</span>
-				<ReactMarkdown
+				<RawMarkdownBlock
+					content={markdownContent}
 					className={`${className}`}
-					rehypePlugins={[rehypeRaw]}
-				>
-					{markdownContent}
-				</ReactMarkdown>
+				/>
 			</div>
 			{markdownContent && finish && iconName === "attempt_completion" &&
 				<div className="infio-chat-message-actions">

src/components/chat-view/Markdown/MermaidBlock.tsx

@@ -1,3 +1,4 @@
+import { CopyIcon } from "lucide-react"
 import mermaid from "mermaid"
 import { memo, useEffect, useRef, useState } from "react"
 import styled from "styled-components"
@@ -76,6 +77,42 @@ interface MermaidBlockProps {
 	code: string
 }
 
+interface MermaidToolbarProps {
+	code: string;
+}
+
+function MermaidToolbar({ code }: MermaidToolbarProps) {
+	const { showCopyFeedback, copyWithFeedback } = useCopyToClipboard()
+
+	const handleCopy = (e: React.MouseEvent) => {
+		e.stopPropagation()
+		// We wrap the code in a markdown block for easy pasting
+		copyWithFeedback("```mermaid\n" + code + "\n```")
+	}
+
+	return (
+		<ToolbarContainer className="mermaid-toolbar">
+			<ToolbarButton className="mermaid-toolbar-btn" onClick={handleCopy} aria-label={t("common:copy_code")}>
+				<CopyIcon size={12} />
+			</ToolbarButton>
+		</ToolbarContainer>
+	)
+}
+
+interface MermaidButtonProps {
+	code: string
+	children: React.ReactNode
+}
+
+function MermaidButton({ code, children }: MermaidButtonProps) {
+	return (
+		<MermaidWrapper>
+			{children}
+			<MermaidToolbar code={code} />
+		</MermaidWrapper>
+	)
+}
+
 function MermaidBlock({ code }: MermaidBlockProps) {
 	const containerRef = useRef<HTMLDivElement>(null)
 	const [isLoading, setIsLoading] = useState(false)
@@ -259,15 +296,78 @@ function MermaidBlock({ code }: MermaidBlockProps) {
 					)}
 				</ErrorContainer>
 			) : (
-				<SvgContainer onClick={handleClick} ref={containerRef} $isLoading={isLoading} />
+				<MermaidButton code={code}>
+					<SvgContainer onClick={handleClick} ref={containerRef} $isLoading={isLoading} />
+				</MermaidButton>
 			)}
 		</MermaidBlockContainer>
 	)
 }
 
-const MermaidBlockContainer = styled.div`
+const MermaidWrapper = styled.div`
 	position: relative;
 	margin: 8px 0;
+
+	&:hover .mermaid-toolbar {
+		opacity: 1;
+	}
+
+	.mermaid-toolbar-btn {
+		display: flex;
+		align-items: center;
+		justify-content: center;
+		background-color: transparent !important;
+		border: none !important;
+		box-shadow: none !important;
+		color: var(--text-muted);
+		padding: 0 !important;
+		margin: 0 !important;
+		width: 24px !important;
+		height: 24px !important;
+
+		&:hover {
+			background-color: var(--background-modifier-hover) !important;
+		}
+	}
+`
+
+const ToolbarContainer = styled.div`
+	position: absolute;
+	top: 8px;
+	right: 8px;
+	z-index: 10;
+	opacity: 0;
+	transition: opacity 0.2s ease-in-out;
+	background-color: var(--background-secondary);
+	border: 1px solid var(--background-modifier-border);
+	border-radius: 6px;
+	padding: 2px;
+	display: flex;
+	align-items: center;
+
+	&:hover {
+		opacity: 1; /* Keep it visible when hovering over the toolbar itself */
+	}
+`
+
+const ToolbarButton = styled.button`
+	padding: 4px;
+	color: var(--text-muted);
+	background: transparent;
+	border: none;
+	cursor: pointer;
+	display: flex;
+	align-items: center;
+	border-radius: 4px;
+
+	&:hover {
+		color: var(--text-normal);
+		background-color: var(--background-modifier-hover);
+	}
+`
+
+const MermaidBlockContainer = styled.div`
+	position: relative;
 `
 
 const LoadingMessage = styled.div`

src/components/chat-view/ReactMarkdown.tsx

@@ -31,8 +31,8 @@ function ReactMarkdown({
 	applyStatus: ApplyStatus
 	onApply: (toolArgs: ToolArgs) => void
 	children: string
-	}) {
-	
+}) {
+
 	const blocks: ParsedMsgBlock[] = useMemo(
 		() => parseMsgBlocks(children),
 		[children],