Skip to content
This repository was archived by the owner on Aug 1, 2025. It is now read-only.

Commit 16f74d0

Browse files
pkukielkapkukielka
committed
Replace rehype-highlight with custom highlighter
1 parent 1716682 commit 16f74d0

3 files changed

Lines changed: 262 additions & 72 deletions

File tree

‎vscode/webviews/components/CustomHJSHighlighter.tsx‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import hljs from 'highlight.js'
2+
import { useMemo } from 'react'
3+
import type React from 'react'
4+
5+
const useHighlightedCode = true
6+
7+
export const CustomHJSHighlighter: React.FC<{
8+
code: string
9+
language?: string
10+
className?: string
11+
}> = ({ code, language, className }) => {
12+
const highlightedCode = useMemo(() => {
13+
if (!code) return ''
14+
15+
try {
16+
// Try to highlight with the specified language
17+
if (language && hljs.getLanguage(language)) {
18+
const result = hljs.highlight(code, { language })
19+
return useHighlightedCode ? result.value : code
20+
}
21+
22+
// Fall back to auto-detection
23+
const result = hljs.highlightAuto(code)
24+
return useHighlightedCode ? result.value : code
25+
} catch (error) {
26+
// If highlighting fails, return plain text
27+
console.warn('Syntax highlighting failed:', error)
28+
return code
29+
}
30+
}, [code, language])
31+
32+
return (
33+
<pre className={`hljs ${className || ''}`}>
34+
<code
35+
className={language ? `language-${language}` : ''}
36+
// biome-ignore lint/security/noDangerouslySetInnerHtml: Required for syntax highlighting
37+
dangerouslySetInnerHTML={{ __html: highlightedCode }}
38+
data-language={language || undefined}
39+
/>
40+
</pre>
41+
)
42+
}
43+
44+
CustomHJSHighlighter.displayName = 'CustomHJSHighlighter'

‎vscode/webviews/components/CustomHTMLHighlighter.tsx‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
/**
2+
* Custom syntax highlighting component that doesn't leak memory
3+
* Replaces rehype-highlight to avoid memory accumulation
4+
*/
5+
6+
import type React from 'react'
7+
import { useMemo } from 'react'
8+
9+
// Simple regex-based syntax highlighting patterns
10+
const HIGHLIGHT_PATTERNS = {
11+
javascript: [
12+
{
13+
pattern: /\b(function|const|let|var|return|if|else|for|while|class|import|export)\b/g,
14+
className: 'hljs-keyword',
15+
},
16+
{ pattern: /\b(true|false|null|undefined)\b/g, className: 'hljs-literal' },
17+
{ pattern: /"([^"\\]|\\.)*"/g, className: 'hljs-string' },
18+
{ pattern: /'([^'\\]|\\.)*'/g, className: 'hljs-string' },
19+
{ pattern: /`([^`\\]|\\.)*`/g, className: 'hljs-string' },
20+
{ pattern: /\/\/.*$/gm, className: 'hljs-comment' },
21+
{ pattern: /\/\*[\s\S]*?\*\//g, className: 'hljs-comment' },
22+
{ pattern: /\b\d+(\.\d+)?\b/g, className: 'hljs-number' },
23+
],
24+
typescript: [
25+
{
26+
pattern:
27+
/\b(function|const|let|var|return|if|else|for|while|class|import|export|interface|type|enum)\b/g,
28+
className: 'hljs-keyword',
29+
},
30+
{ pattern: /\b(string|number|boolean|any|void|never)\b/g, className: 'hljs-type' },
31+
{ pattern: /\b(true|false|null|undefined)\b/g, className: 'hljs-literal' },
32+
{ pattern: /"([^"\\]|\\.)*"/g, className: 'hljs-string' },
33+
{ pattern: /'([^'\\]|\\.)*'/g, className: 'hljs-string' },
34+
{ pattern: /`([^`\\]|\\.)*`/g, className: 'hljs-string' },
35+
{ pattern: /\/\/.*$/gm, className: 'hljs-comment' },
36+
{ pattern: /\/\*[\s\S]*?\*\//g, className: 'hljs-comment' },
37+
{ pattern: /\b\d+(\.\d+)?\b/g, className: 'hljs-number' },
38+
],
39+
python: [
40+
{
41+
pattern: /\b(def|class|import|from|return|if|else|elif|for|while|try|except|with|as)\b/g,
42+
className: 'hljs-keyword',
43+
},
44+
{ pattern: /\b(True|False|None)\b/g, className: 'hljs-literal' },
45+
{ pattern: /"([^"\\]|\\.)*"/g, className: 'hljs-string' },
46+
{ pattern: /'([^'\\]|\\.)*'/g, className: 'hljs-string' },
47+
{ pattern: /#.*$/gm, className: 'hljs-comment' },
48+
{ pattern: /\b\d+(\.\d+)?\b/g, className: 'hljs-number' },
49+
],
50+
go: [
51+
{
52+
pattern: /\b(func|var|const|type|import|package|return|if|else|for|range|switch|case)\b/g,
53+
className: 'hljs-keyword',
54+
},
55+
{ pattern: /\b(true|false|nil)\b/g, className: 'hljs-literal' },
56+
{ pattern: /"([^"\\]|\\.)*"/g, className: 'hljs-string' },
57+
{ pattern: /`([^`\\]|\\.)*`/g, className: 'hljs-string' },
58+
{ pattern: /\/\/.*$/gm, className: 'hljs-comment' },
59+
{ pattern: /\/\*[\s\S]*?\*\//g, className: 'hljs-comment' },
60+
{ pattern: /\b\d+(\.\d+)?\b/g, className: 'hljs-number' },
61+
],
62+
}
63+
64+
const DEFAULT_PATTERNS = [
65+
{ pattern: /"([^"\\]|\\.)*"/g, className: 'hljs-string' },
66+
{ pattern: /'([^'\\]|\\.)*'/g, className: 'hljs-string' },
67+
{ pattern: /\b\d+(\.\d+)?\b/g, className: 'hljs-number' },
68+
]
69+
70+
function highlightCode(code: string, language?: string): string {
71+
if (!code) return ''
72+
73+
const patterns = HIGHLIGHT_PATTERNS[language as keyof typeof HIGHLIGHT_PATTERNS] || DEFAULT_PATTERNS
74+
75+
// Escape HTML first
76+
const escaped = code.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;')
77+
78+
// Find all matches for all patterns
79+
const matches: Array<{ start: number; end: number; className: string; text: string }> = []
80+
81+
for (const { pattern, className } of patterns) {
82+
pattern.lastIndex = 0 // Reset regex
83+
let match = pattern.exec(escaped)
84+
while (match !== null) {
85+
matches.push({
86+
start: match.index,
87+
end: match.index + match[0].length,
88+
className,
89+
text: match[0],
90+
})
91+
if (!pattern.global) break
92+
match = pattern.exec(escaped)
93+
}
94+
}
95+
96+
// Sort matches by position and filter out overlaps
97+
matches.sort((a, b) => a.start - b.start)
98+
const filteredMatches = []
99+
let lastEnd = 0
100+
101+
for (const match of matches) {
102+
if (match.start >= lastEnd) {
103+
filteredMatches.push(match)
104+
lastEnd = match.end
105+
}
106+
}
107+
108+
// Build the final highlighted string
109+
let result = ''
110+
let currentIndex = 0
111+
112+
for (const match of filteredMatches) {
113+
// Add text before match
114+
if (match.start > currentIndex) {
115+
result += escaped.slice(currentIndex, match.start)
116+
}
117+
// Add highlighted match
118+
result += `<span class="${match.className}">${match.text}</span>`
119+
currentIndex = match.end
120+
}
121+
122+
// Add remaining text
123+
if (currentIndex < escaped.length) {
124+
result += escaped.slice(currentIndex)
125+
}
126+
127+
return result
128+
}
129+
130+
export const CustomHTMLHighlighter: React.FC<{
131+
code: string
132+
language?: string
133+
className?: string
134+
}> = ({ code, language, className }) => {
135+
const highlightedCode = useMemo(() => {
136+
return highlightCode(code, language)
137+
}, [code, language])
138+
139+
return (
140+
<pre className={`hljs ${className || ''}`}>
141+
<code
142+
className={language ? `language-${language}` : ''}
143+
// biome-ignore lint/security/noDangerouslySetInnerHtml: <explanation>
144+
dangerouslySetInnerHTML={{ __html: highlightedCode }}
145+
/>
146+
</pre>
147+
)
148+
}

‎vscode/webviews/components/MarkdownFromCody.tsx‎

Lines changed: 70 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
import { CodyIDE } from '@sourcegraph/cody-shared'
22
import type { ComponentProps, FunctionComponent } from 'react'
3-
import { useMemo } from 'react'
3+
import { memo, useMemo } from 'react'
44
import Markdown, { defaultUrlTransform } from 'react-markdown'
55
import type { Components, UrlTransform } from 'react-markdown/lib'
6-
import rehypeHighlight, { type Options as RehypeHighlightOptions } from 'rehype-highlight'
6+
77
import rehypeSanitize, { type Options as RehypeSanitizeOptions, defaultSchema } from 'rehype-sanitize'
88
import remarkGFM from 'remark-gfm'
99
import type { Pluggable } from 'unified/lib'
1010
import { remarkAttachFilePathToCodeBlocks } from '../chat/extract-file-path'
1111
import { SYNTAX_HIGHLIGHTING_LANGUAGES } from '../utils/highlight'
1212
import { useConfig } from '../utils/useConfig'
13+
import { CustomHJSHighlighter } from './CustomHJSHighlighter'
1314

1415
/**
1516
* Supported URIs to render as links in outputted markdown.
@@ -84,7 +85,7 @@ const URL_PROCESSORS: Partial<Record<CodyIDE, UrlTransform>> = {
8485
}
8586

8687
/**
87-
* Transforms the children string by wrapping it in one extra backtick if we find '```markdown'.
88+
* Transforms the children string by wrapping it in one extra backtick if we find '````markdown'.
8889
* This is used to preserve the formatting of Markdown code blocks within the Markdown content.
8990
* Such cases happen when you ask Cody to create a Markdown file or when you load a history chat
9091
* that contains replies for creating Markdown files.
@@ -103,83 +104,80 @@ const childrenTransform = (children: string): string => {
103104
return children.slice(0, lastIdx) + '````' + children.slice(lastIdx + 3)
104105
}
105106

107+
const _markdownPluginProps: Pick<ComponentProps<typeof Markdown>, 'rehypePlugins' | 'remarkPlugins'> = {
108+
rehypePlugins: [
109+
[
110+
rehypeSanitize,
111+
{
112+
...defaultSchema,
113+
tagNames: ALLOWED_ELEMENTS,
114+
attributes: {
115+
...defaultSchema.attributes,
116+
code: [
117+
...(defaultSchema.attributes?.code || []),
118+
// Allow various metadata attributes for code blocks
119+
['data-file-path'],
120+
['data-is-code-complete'],
121+
['data-language'],
122+
['data-source-text'],
123+
[
124+
'className',
125+
...Object.keys(SYNTAX_HIGHLIGHTING_LANGUAGES).map(
126+
language => `language-${language}`
127+
),
128+
],
129+
],
130+
},
131+
} satisfies RehypeSanitizeOptions,
132+
],
133+
],
134+
remarkPlugins: [remarkGFM, remarkAttachFilePathToCodeBlocks],
135+
}
136+
106137
export const MarkdownFromCody: FunctionComponent<{
107138
className?: string
108139
prefixRemarkPlugins?: Pluggable[]
109140
components?: Partial<Components>
110141
children: string
111-
}> = ({ className, prefixRemarkPlugins, components, children }) => {
142+
}> = memo(({ className, children, components }) => {
112143
const clientType = useConfig().clientCapabilities.agentIDE
113144
const urlTransform = useMemo(() => URL_PROCESSORS[clientType] ?? defaultUrlProcessor, [clientType])
114-
const chatReplyTransformed = childrenTransform(children)
145+
const chatReplyTransformed = useMemo(() => childrenTransform(children), [children])
115146

116-
return (
117-
<Markdown
118-
className={className}
119-
{...markdownPluginProps(prefixRemarkPlugins ?? [])}
120-
urlTransform={urlTransform}
121-
components={components ?? {}}
122-
>
123-
{chatReplyTransformed}
124-
</Markdown>
147+
const markdownComponents = useMemo(
148+
() =>
149+
({
150+
...components,
151+
code: ({ node, className, children, ...props }: any) => {
152+
const match = /language-(\w+)/.exec(className || '')
153+
const language = match ? match[1] : undefined
154+
const code = String(children).replace(/\n$/, '')
155+
156+
if (!language) {
157+
return (
158+
<code className={className} {...props}>
159+
{children}
160+
</code>
161+
)
162+
}
163+
164+
return <CustomHJSHighlighter code={code} language={language} className={className} />
165+
},
166+
}) as Partial<Components>,
167+
[components]
125168
)
126-
}
127169

128-
let _markdownPluginProps: ReturnType<typeof markdownPluginProps> | undefined
129-
function markdownPluginProps(
130-
prefixRemarkPlugins: Pluggable[] = []
131-
): Pick<ComponentProps<typeof Markdown>, 'rehypePlugins' | 'remarkPlugins'> {
132-
if (_markdownPluginProps) {
133-
return _markdownPluginProps
134-
}
170+
return (
171+
<div className={className}>
172+
<Markdown
173+
{..._markdownPluginProps}
174+
urlTransform={urlTransform}
175+
components={markdownComponents}
176+
>
177+
{chatReplyTransformed}
178+
</Markdown>
179+
</div>
180+
)
181+
})
135182

136-
_markdownPluginProps = {
137-
rehypePlugins: [
138-
[
139-
rehypeSanitize,
140-
{
141-
...defaultSchema,
142-
tagNames: ALLOWED_ELEMENTS,
143-
attributes: {
144-
...defaultSchema.attributes,
145-
code: [
146-
...(defaultSchema.attributes?.code || []),
147-
// Allow various metadata attributes for code blocks
148-
['data-file-path'],
149-
['data-is-code-complete'],
150-
['data-language'],
151-
['data-source-text'],
152-
[
153-
'className',
154-
...Object.keys(SYNTAX_HIGHLIGHTING_LANGUAGES).map(
155-
language => `language-${language}`
156-
),
157-
],
158-
],
159-
},
160-
} satisfies RehypeSanitizeOptions,
161-
],
162-
[
163-
// HACK(sqs): Need to use rehype-highlight@^6.0.0 to avoid a memory leak
164-
// (https://github.com/remarkjs/react-markdown/issues/791), but the types are
165-
// slightly off.
166-
rehypeHighlight as any,
167-
{
168-
detect: true,
169-
languages: {
170-
...SYNTAX_HIGHLIGHTING_LANGUAGES,
171-
},
172-
173-
// `ignoreMissing: true` is required to avoid errors when trying to highlight
174-
// partial code blocks received from the LLM that have (e.g.) "```p" for
175-
// "```python". This is only needed on rehype-highlight@^6.0.0, which we needed
176-
// to downgrade to in order to avoid a memory leak
177-
// (https://github.com/remarkjs/react-markdown/issues/791).
178-
ignoreMissing: true,
179-
} satisfies RehypeHighlightOptions & { ignoreMissing: boolean },
180-
],
181-
],
182-
remarkPlugins: [...prefixRemarkPlugins, remarkGFM, remarkAttachFilePathToCodeBlocks],
183-
}
184-
return _markdownPluginProps
185-
}
183+
MarkdownFromCody.displayName = 'MarkdownFromCody'

0 commit comments

Comments
 (0)