Skip to content

Commit 10371ad

Browse files
daryllimytdaryllimyt
authored
perf(ui): reduce chat session CPU usage from tool call re-renders (#2452)
1 parent 571b6bf commit 10371ad

File tree

4 files changed

+597
-140
lines changed

4 files changed

+597
-140
lines changed

‎frontend/src/components/ai-elements/tool.tsx‎

Lines changed: 119 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
import {
1313
type ComponentProps,
1414
isValidElement,
15+
memo,
1516
type ReactNode,
1617
useEffect,
1718
useMemo,
@@ -283,134 +284,143 @@ function formatBytes(bytes: number): string {
283284
return `${mb.toFixed(1)} MB`
284285
}
285286

286-
function ToolPayload({
287-
payload,
288-
downloadName,
289-
}: {
290-
payload: unknown
291-
downloadName: string
292-
}) {
293-
const normalizedPayload = useMemo(
294-
() => extractMcpTextContent(payload),
295-
[payload]
296-
)
297-
const serialized = useMemo(
298-
() => serializeToolPayload(normalizedPayload),
299-
[normalizedPayload]
300-
)
301-
const isLargePayload = serialized.text.length > MAX_INLINE_PAYLOAD_CHARS
302-
const codeLanguage = serialized.extension === "json" ? "json" : "console"
303-
const byteCount = useMemo(
304-
() =>
305-
isLargePayload ? new TextEncoder().encode(serialized.text).length : 0,
306-
[isLargePayload, serialized.text]
307-
)
308-
const [downloadHref, setDownloadHref] = useState<string | null>(null)
287+
const ToolPayload = memo(
288+
function ToolPayload({
289+
payload,
290+
downloadName,
291+
}: {
292+
payload: unknown
293+
downloadName: string
294+
}) {
295+
const normalizedPayload = useMemo(
296+
() => extractMcpTextContent(payload),
297+
[payload]
298+
)
299+
const serialized = useMemo(
300+
() => serializeToolPayload(normalizedPayload),
301+
[normalizedPayload]
302+
)
303+
const isLargePayload = serialized.text.length > MAX_INLINE_PAYLOAD_CHARS
304+
const codeLanguage = serialized.extension === "json" ? "json" : "console"
305+
const byteCount = useMemo(
306+
() =>
307+
isLargePayload ? new TextEncoder().encode(serialized.text).length : 0,
308+
[isLargePayload, serialized.text]
309+
)
310+
const [downloadHref, setDownloadHref] = useState<string | null>(null)
309311

310-
useEffect(() => {
311-
if (!isLargePayload || !serialized.text) {
312-
setDownloadHref(null)
313-
return
314-
}
312+
useEffect(() => {
313+
if (!isLargePayload || !serialized.text) {
314+
setDownloadHref(null)
315+
return
316+
}
315317

316-
const blob = new Blob([serialized.text], {
317-
type: serialized.extension === "json" ? "application/json" : "text/plain",
318-
})
319-
const href = URL.createObjectURL(blob)
320-
setDownloadHref(href)
321-
return () => URL.revokeObjectURL(href)
322-
}, [isLargePayload, serialized.extension, serialized.text])
318+
const blob = new Blob([serialized.text], {
319+
type:
320+
serialized.extension === "json" ? "application/json" : "text/plain",
321+
})
322+
const href = URL.createObjectURL(blob)
323+
setDownloadHref(href)
324+
return () => URL.revokeObjectURL(href)
325+
}, [isLargePayload, serialized.extension, serialized.text])
326+
327+
if (!serialized.text) {
328+
return null
329+
}
323330

324-
if (!serialized.text) {
325-
return null
326-
}
331+
if (isLargePayload) {
332+
return (
333+
<div className="space-y-2">
334+
<div className="flex items-center justify-between gap-2 rounded-md border border-dashed bg-background/70 px-2.5 py-2 text-xs text-muted-foreground">
335+
<span>
336+
Large payload ({formatBytes(byteCount)}). Preview hidden.
337+
</span>
338+
{downloadHref && (
339+
<Button
340+
asChild
341+
variant="outline"
342+
size="sm"
343+
className="h-6 px-2 text-xs"
344+
>
345+
<a
346+
href={downloadHref}
347+
download={`${downloadName}.${serialized.extension}`}
348+
>
349+
<DownloadIcon className="mr-1.5 size-3" />
350+
Download file
351+
</a>
352+
</Button>
353+
)}
354+
</div>
355+
</div>
356+
)
357+
}
327358

328-
if (isLargePayload) {
329359
return (
330360
<div className="space-y-2">
331-
<div className="flex items-center justify-between gap-2 rounded-md border border-dashed bg-background/70 px-2.5 py-2 text-xs text-muted-foreground">
332-
<span>Large payload ({formatBytes(byteCount)}). Preview hidden.</span>
333-
{downloadHref && (
334-
<Button
335-
asChild
336-
variant="outline"
337-
size="sm"
338-
className="h-6 px-2 text-xs"
339-
>
340-
<a
341-
href={downloadHref}
342-
download={`${downloadName}.${serialized.extension}`}
343-
>
344-
<DownloadIcon className="mr-1.5 size-3" />
345-
Download file
346-
</a>
347-
</Button>
348-
)}
349-
</div>
361+
<CodeBlock
362+
code={serialized.text}
363+
language={codeLanguage}
364+
className="[&_code]:text-[11px] [&_pre]:px-3 [&_pre]:py-2.5 [&_pre]:text-[11px]"
365+
>
366+
<CodeBlockCopyButton className="absolute right-1.5 top-1.5 z-10 size-5 [&_svg]:size-3" />
367+
</CodeBlock>
350368
</div>
351369
)
352-
}
353-
354-
return (
355-
<div className="space-y-2">
356-
<CodeBlock
357-
code={serialized.text}
358-
language={codeLanguage}
359-
className="[&_code]:text-[11px] [&_pre]:px-3 [&_pre]:py-2.5 [&_pre]:text-[11px]"
360-
>
361-
<CodeBlockCopyButton className="absolute right-1.5 top-1.5 z-10 size-5 [&_svg]:size-3" />
362-
</CodeBlock>
363-
</div>
364-
)
365-
}
370+
},
371+
(prev, next) =>
372+
prev.payload === next.payload && prev.downloadName === next.downloadName
373+
)
374+
ToolPayload.displayName = "ToolPayload"
366375

367376
export type ToolInputProps = ComponentProps<"div"> & {
368377
input: ToolPart["input"]
369378
}
370379

371-
export const ToolInput = ({ className, input, ...props }: ToolInputProps) => (
372-
<div className={cn("space-y-2 overflow-hidden", className)} {...props}>
373-
<h4 className="font-medium text-[11px] text-muted-foreground tracking-wide">
374-
PARAMETERS
375-
</h4>
376-
<ToolPayload payload={input} downloadName="tool-parameters" />
377-
</div>
380+
export const ToolInput = memo(
381+
({ className, input, ...props }: ToolInputProps) => (
382+
<div className={cn("space-y-2 overflow-hidden", className)} {...props}>
383+
<h4 className="font-medium text-[11px] text-muted-foreground tracking-wide">
384+
PARAMETERS
385+
</h4>
386+
<ToolPayload payload={input} downloadName="tool-parameters" />
387+
</div>
388+
)
378389
)
390+
ToolInput.displayName = "ToolInput"
379391

380392
export type ToolOutputProps = ComponentProps<"div"> & {
381393
output: ToolPart["output"]
382394
errorText: ToolPart["errorText"]
383395
}
384396

385-
export const ToolOutput = ({
386-
className,
387-
output,
388-
errorText,
389-
...props
390-
}: ToolOutputProps) => {
391-
const hasOutput = output !== undefined && output !== null
392-
if (!hasOutput && !errorText) {
393-
return null
394-
}
397+
export const ToolOutput = memo(
398+
({ className, output, errorText, ...props }: ToolOutputProps) => {
399+
const hasOutput = output !== undefined && output !== null
400+
if (!hasOutput && !errorText) {
401+
return null
402+
}
395403

396-
return (
397-
<div className={cn("space-y-2", className)} {...props}>
398-
<h4 className="font-medium text-[11px] text-muted-foreground tracking-wide">
399-
{errorText ? "Error" : "RESULT"}
400-
</h4>
401-
{errorText && (
402-
<div className="rounded-md bg-destructive/10 px-2.5 py-1.5 text-[11px] text-destructive">
403-
{errorText}
404-
</div>
405-
)}
406-
{hasOutput &&
407-
(isValidElement(output) ? (
408-
<div className="rounded-md bg-muted/50 p-1.5 text-[11px] text-foreground">
409-
{output}
404+
return (
405+
<div className={cn("space-y-2", className)} {...props}>
406+
<h4 className="font-medium text-[11px] text-muted-foreground tracking-wide">
407+
{errorText ? "Error" : "RESULT"}
408+
</h4>
409+
{errorText && (
410+
<div className="rounded-md bg-destructive/10 px-2.5 py-1.5 text-[11px] text-destructive">
411+
{errorText}
410412
</div>
411-
) : (
412-
<ToolPayload payload={output} downloadName="tool-result" />
413-
))}
414-
</div>
415-
)
416-
}
413+
)}
414+
{hasOutput &&
415+
(isValidElement(output) ? (
416+
<div className="rounded-md bg-muted/50 p-1.5 text-[11px] text-foreground">
417+
{output}
418+
</div>
419+
) : (
420+
<ToolPayload payload={output} downloadName="tool-result" />
421+
))}
422+
</div>
423+
)
424+
}
425+
)
426+
ToolOutput.displayName = "ToolOutput"

0 commit comments

Comments
 (0)