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

Commit bad1551

Browse files
pkukielkapkukielka
authored
Improve markdown component caching (#8183)
This PR is not fixing a memory leak in the DOM highlighter code but significantly reduces amount of re-computation which is done. With mocked LLM response of length of 7528 characters I get: **Before my changes:** Execution time: 1 min 25 s ec Peak memory used: over 2GB Stable memory after final GC: 833MB **After my changes:** Execution time: **45 sec** Peak memory used: over 1GB Stable memory after final GC: 719MB Sop while memory usage drop was not significant there is very visible speedup in the execution speed. Also memory usage swings were subjectively less significant. ## Test plan I do not have good testing procedure on `main`, but I di testing using my testing PR, with some custom logging: https://github.com/sourcegraph/cody/pull/8167
1 parent bcf9252 commit bad1551

3 files changed

Lines changed: 175 additions & 148 deletions

File tree

‎vscode/webviews/chat/ChatMessageContent/ChatMessageContent.tsx‎

Lines changed: 70 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import type { Guardrails, PromptString } from '@sourcegraph/cody-shared'
22
import { clsx } from 'clsx'
33
import type React from 'react'
4-
import { useCallback, useMemo } from 'react'
4+
import { memo, useCallback, useMemo } from 'react'
55
import { RichMarkdown } from '../../components/RichMarkdown'
66
import { getVSCodeAPI } from '../../utils/VSCodeApi'
77
import { useConfig } from '../../utils/useConfig'
@@ -51,71 +51,79 @@ interface ChatMessageContentProps {
5151
/**
5252
* A component presenting the content of a chat message.
5353
*/
54-
export const ChatMessageContent: React.FunctionComponent<ChatMessageContentProps> = ({
55-
displayMarkdown,
56-
isMessageLoading,
57-
humanMessage,
58-
copyButtonOnSubmit,
59-
insertButtonOnSubmit,
60-
onRegenerate,
61-
regeneratingCodeBlocks,
62-
guardrails,
63-
className,
64-
smartApply,
65-
isThoughtProcessOpened,
66-
setThoughtProcessOpened,
67-
}) => {
68-
const config = useConfig()
54+
export const ChatMessageContent: React.FunctionComponent<ChatMessageContentProps> = memo(
55+
({
56+
displayMarkdown,
57+
isMessageLoading,
58+
humanMessage,
59+
copyButtonOnSubmit,
60+
insertButtonOnSubmit,
61+
onRegenerate,
62+
regeneratingCodeBlocks,
63+
guardrails,
64+
className,
65+
smartApply,
66+
isThoughtProcessOpened,
67+
setThoughtProcessOpened,
68+
}) => {
69+
const config = useConfig()
6970

70-
const { displayContent, thinkContent, isThinking } = useMemo(
71-
() => extractThinkContent(displayMarkdown),
72-
[displayMarkdown]
73-
)
71+
const { displayContent, thinkContent, isThinking } = useMemo(
72+
() => extractThinkContent(displayMarkdown),
73+
[displayMarkdown]
74+
)
7475

75-
const onInsert = config.config.hasEditCapability ? insertButtonOnSubmit : undefined
76+
const onInsert = useMemo(
77+
() => (config.config.hasEditCapability ? insertButtonOnSubmit : undefined),
78+
[config.config.hasEditCapability, insertButtonOnSubmit]
79+
)
7680

77-
let onExecute: ((command: string) => void) | undefined = useCallback((command: string) => {
78-
// Execute command in terminal
79-
const vscodeAPI = getVSCodeAPI()
80-
vscodeAPI.postMessage({
81-
command: 'command',
82-
id: 'cody.terminal.execute',
83-
arg: command.trim(),
84-
})
85-
}, [])
81+
let onExecute: ((command: string) => void) | undefined = useCallback((command: string) => {
82+
// Execute command in terminal
83+
const vscodeAPI = getVSCodeAPI()
84+
vscodeAPI.postMessage({
85+
command: 'command',
86+
id: 'cody.terminal.execute',
87+
arg: command.trim(),
88+
})
89+
}, [])
8690

87-
// TODO: Replace this isVSCode check with a client capability check for
88-
// terminal execution when agent/src/vscode-shim.ts implements `terminal()`
89-
onExecute = config.clientCapabilities.isVSCode ? onExecute : undefined
91+
// TODO: Replace this isVSCode check with a client capability check for
92+
// terminal execution when agent/src/vscode-shim.ts implements `terminal()`
93+
onExecute = useMemo(
94+
() => (config.clientCapabilities.isVSCode ? onExecute : undefined),
95+
[config.clientCapabilities.isVSCode, onExecute]
96+
)
9097

91-
const onCopy = useCallback(
92-
(code: string) => copyButtonOnSubmit?.(code, 'Button'),
93-
[copyButtonOnSubmit]
94-
)
98+
const onCopy = useCallback(
99+
(code: string) => copyButtonOnSubmit?.(code, 'Button'),
100+
[copyButtonOnSubmit]
101+
)
95102

96-
return (
97-
<div data-testid="chat-message-content">
98-
{setThoughtProcessOpened && thinkContent.length > 0 && (
99-
<ThinkingCell
100-
isOpen={!!isThoughtProcessOpened}
101-
setIsOpen={setThoughtProcessOpened}
102-
isThinking={isMessageLoading && isThinking}
103-
thought={thinkContent}
103+
return (
104+
<div data-testid="chat-message-content">
105+
{setThoughtProcessOpened && thinkContent.length > 0 && (
106+
<ThinkingCell
107+
isOpen={!!isThoughtProcessOpened}
108+
setIsOpen={setThoughtProcessOpened}
109+
isThinking={isMessageLoading && isThinking}
110+
thought={thinkContent}
111+
/>
112+
)}
113+
<RichMarkdown
114+
markdown={displayContent}
115+
isMessageLoading={isMessageLoading}
116+
guardrails={guardrails}
117+
onCopy={onCopy}
118+
onInsert={onInsert}
119+
onExecute={onExecute}
120+
onRegenerate={onRegenerate}
121+
regeneratingCodeBlocks={regeneratingCodeBlocks}
122+
smartApply={smartApply}
123+
className={clsx(styles.content, className)}
124+
hasEditIntent={humanMessage?.intent === 'edit'}
104125
/>
105-
)}
106-
<RichMarkdown
107-
markdown={displayContent}
108-
isMessageLoading={isMessageLoading}
109-
guardrails={guardrails}
110-
onCopy={onCopy}
111-
onInsert={onInsert}
112-
onExecute={onExecute}
113-
onRegenerate={onRegenerate}
114-
regeneratingCodeBlocks={regeneratingCodeBlocks}
115-
smartApply={smartApply}
116-
className={clsx(styles.content, className)}
117-
hasEditIntent={humanMessage?.intent === 'edit'}
118-
/>
119-
</div>
120-
)
121-
}
126+
</div>
127+
)
128+
}
129+
)

‎vscode/webviews/chat/Transcript.tsx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -687,7 +687,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
687687

688688
const onRegenerate = useCallback(
689689
(code: string, language?: string) => {
690-
if (assistantMessage) {
690+
if (assistantMessage?.index) {
691691
const id = uuid.v4()
692692
regenerateCodeBlock({ id, code, language, index: assistantMessage.index })
693693
setRegeneratingCodeBlocks(blocks => [
@@ -698,7 +698,7 @@ const TranscriptInteraction: FC<TranscriptInteractionProps> = memo(props => {
698698
console.warn('tried to regenerate a code block, but there is no assistant message')
699699
}
700700
},
701-
[assistantMessage]
701+
[assistantMessage?.index]
702702
)
703703

704704
const isAgenticMode = useMemo(

‎vscode/webviews/components/RichMarkdown.tsx‎

Lines changed: 103 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import type { Guardrails } from '@sourcegraph/cody-shared'
22
import { clsx } from 'clsx'
33
import type { Code, Root } from 'mdast'
44
import type React from 'react'
5+
import { memo, useCallback, useMemo } from 'react'
56
import type { Plugin } from 'unified'
67
import { visit } from 'unist-util-visit'
78
import type { CodeBlockActionsProps } from '../chat/ChatMessageContent/ChatMessageContent'
@@ -70,96 +71,114 @@ export const remarkAttachCompletedCodeBlocks: Plugin<[], Root> = () => {
7071
* which provides syntax highlighting, action buttons, and optional guardrails
7172
* protection.
7273
*/
73-
export const RichMarkdown: React.FC<RichMarkdownProps> = ({
74-
markdown,
75-
isMessageLoading,
76-
regeneratingCodeBlocks,
77-
guardrails,
78-
onCopy,
79-
onInsert,
80-
onExecute,
81-
onRegenerate,
82-
smartApply,
83-
className,
84-
hasEditIntent,
85-
}) => {
86-
// Handle rendering of code blocks with our custom RichCodeBlock component
87-
const components = {
88-
pre({ node, inline, className, children, ...props }: any) {
89-
// Don't process inline code blocks
90-
if (inline) {
91-
return (
92-
<code className={className} {...props}>
93-
{children}
94-
</code>
95-
)
74+
export const RichMarkdown: React.FC<RichMarkdownProps> = memo(
75+
({
76+
markdown,
77+
isMessageLoading,
78+
regeneratingCodeBlocks,
79+
guardrails,
80+
onCopy,
81+
onInsert,
82+
onExecute,
83+
onRegenerate,
84+
smartApply,
85+
className,
86+
hasEditIntent,
87+
}) => {
88+
// Memoize the extractText function to avoid recreating it every render
89+
const extractText = useCallback((node: any): string => {
90+
if (typeof node === 'string') return node
91+
if (!node) return ''
92+
if (node.type === 'text' && node.value) return node.value
93+
if (node.children) {
94+
return node.children.map(extractText).join('')
9695
}
96+
return ''
97+
}, [])
9798

98-
// Get the code node (if it exists)
99-
const codeNode =
100-
node.children.length === 1 && node.children[0].type === 'element'
101-
? node.children[0]
102-
: null
99+
// Handle rendering of code blocks with our custom RichCodeBlock component
100+
const components = useMemo(
101+
() => ({
102+
pre({ node, inline, className, children, ...props }: any) {
103+
// Don't process inline code blocks
104+
if (inline) {
105+
return (
106+
<code className={className} {...props}>
107+
{children}
108+
</code>
109+
)
110+
}
103111

104-
// Get the cached highlighting result (if there is a key, and if the result is cached)
105-
const {
106-
'data-source-text': sourceText,
107-
'data-is-code-complete': isThisBlockComplete,
108-
'data-file-path': filePath,
109-
'data-language': language,
110-
} = (codeNode?.properties as TerminatedCodeData['hProperties'] | undefined) || {
111-
'data-is-code-complete': false,
112-
}
112+
// Get the code node (if it exists)
113+
const codeNode =
114+
node.children.length === 1 && node.children[0].type === 'element'
115+
? node.children[0]
116+
: null
113117

114-
const extractText = (node: any): string => {
115-
if (typeof node === 'string') return node
116-
if (!node) return ''
117-
if (node.type === 'text' && node.value) return node.value
118-
if (node.children) {
119-
return node.children.map(extractText).join('')
120-
}
121-
return ''
122-
}
123-
const plainText = extractText(node)
118+
// Get the cached highlighting result (if there is a key, and if the result is cached)
119+
const {
120+
'data-source-text': sourceText,
121+
'data-is-code-complete': isThisBlockComplete,
122+
'data-file-path': filePath,
123+
'data-language': language,
124+
} = (codeNode?.properties as TerminatedCodeData['hProperties'] | undefined) || {
125+
'data-is-code-complete': false,
126+
}
124127

125-
// Determine if this is a shell command
126-
const isShellCommand = language === 'bash' || language === 'sh'
128+
const plainText = extractText(node)
127129

128-
const regenerating = regeneratingCodeBlocks.find(
129-
block => block.code === plainText && !block.error
130-
)
130+
// Determine if this is a shell command
131+
const isShellCommand = language === 'bash' || language === 'sh'
131132

132-
// Render with our RichCodeBlock component
133-
return (
134-
<RichCodeBlock
135-
hasEditIntent={hasEditIntent}
136-
plainCode={plainText}
137-
markdownCode={sourceText ?? ''}
138-
language={language}
139-
fileName={filePath}
140-
isMessageLoading={isMessageLoading}
141-
isCodeComplete={!regenerating && (isThisBlockComplete || !isMessageLoading)}
142-
guardrails={guardrails}
143-
onCopy={onCopy}
144-
onInsert={onInsert}
145-
onExecute={isShellCommand ? onExecute : undefined}
146-
onRegenerate={onRegenerate}
147-
smartApply={smartApply}
133+
const regenerating = regeneratingCodeBlocks.find(
134+
block => block.code === plainText && !block.error
135+
)
136+
137+
// Render with our RichCodeBlock component
138+
return (
139+
<RichCodeBlock
140+
hasEditIntent={hasEditIntent}
141+
plainCode={plainText}
142+
markdownCode={sourceText ?? ''}
143+
language={language}
144+
fileName={filePath}
145+
isMessageLoading={isMessageLoading}
146+
isCodeComplete={!regenerating && (isThisBlockComplete || !isMessageLoading)}
147+
guardrails={guardrails}
148+
onCopy={onCopy}
149+
onInsert={onInsert}
150+
onExecute={isShellCommand ? onExecute : undefined}
151+
onRegenerate={onRegenerate}
152+
smartApply={smartApply}
153+
>
154+
{children}
155+
</RichCodeBlock>
156+
)
157+
},
158+
}),
159+
[
160+
extractText,
161+
hasEditIntent,
162+
regeneratingCodeBlocks,
163+
isMessageLoading,
164+
guardrails,
165+
onCopy,
166+
onInsert,
167+
onExecute,
168+
onRegenerate,
169+
smartApply,
170+
]
171+
)
172+
173+
return (
174+
<div className={clsx('markdown-content', className)}>
175+
<MarkdownFromCody
176+
components={components}
177+
prefixRemarkPlugins={[remarkAttachCompletedCodeBlocks]}
148178
>
149-
{children}
150-
</RichCodeBlock>
151-
)
152-
},
179+
{markdown}
180+
</MarkdownFromCody>
181+
</div>
182+
)
153183
}
154-
155-
return (
156-
<div className={clsx('markdown-content', className)}>
157-
<MarkdownFromCody
158-
components={components}
159-
prefixRemarkPlugins={[remarkAttachCompletedCodeBlocks]}
160-
>
161-
{markdown}
162-
</MarkdownFromCody>
163-
</div>
164-
)
165-
}
184+
)

0 commit comments

Comments
 (0)