Skip to content

Commit 472e7c6

Browse files
jahoomajahooma
committed
Add back memo() and don't prop drill streamingAgents
1 parent e1e02df commit 472e7c6

File tree

8 files changed

+49
-74
lines changed

8 files changed

+49
-74
lines changed

‎cli/src/components/blocks/agent-block-grid.tsx‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ export interface AgentBlockGridProps {
77
agentBlocks: AgentContentBlock[]
88
keyPrefix: string
99
availableWidth: number
10-
streamingAgents: Set<string>
1110
renderAgentBranch: (
1211
agentBlock: AgentContentBlock,
1312
keyPrefix: string,

‎cli/src/components/blocks/agent-branch-wrapper.tsx‎

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { ContentWithMarkdown } from './content-with-markdown'
99
import { ThinkingBlock } from './thinking-block'
1010
import { trimTrailingNewlines, sanitizePreview } from './block-helpers'
1111
import { useTheme } from '../../hooks/use-theme'
12+
import { useChatStore } from '../../state/chat-store'
1213
import { AGENT_CONTENT_HORIZONTAL_PADDING } from '../../utils/layout-helpers'
1314
import { shouldRenderAsSimpleText } from '../../utils/constants'
1415
import { isImplementorAgent, getImplementorIndex } from '../../utils/implementor-helpers'
@@ -30,7 +31,6 @@ interface AgentBodyProps {
3031
parentIsStreaming: boolean
3132
availableWidth: number
3233
markdownPalette: MarkdownPalette
33-
streamingAgents: Set<string>
3434
onToggleCollapsed: (id: string) => void
3535
onBuildFast: () => void
3636
onBuildMax: () => void
@@ -44,7 +44,6 @@ interface AgentBodyPropsRef {
4444
parentIsStreaming: boolean
4545
availableWidth: number
4646
markdownPalette: MarkdownPalette
47-
streamingAgents: Set<string>
4847
onToggleCollapsed: (id: string) => void
4948
onBuildFast: () => void
5049
onBuildMax: () => void
@@ -60,7 +59,6 @@ const AgentBody = memo(
6059
parentIsStreaming,
6160
availableWidth,
6261
markdownPalette,
63-
streamingAgents,
6462
onToggleCollapsed,
6563
onBuildFast,
6664
onBuildMax,
@@ -94,7 +92,6 @@ const AgentBody = memo(
9492
parentIsStreaming,
9593
availableWidth,
9694
markdownPalette,
97-
streamingAgents,
9895
onToggleCollapsed,
9996
onBuildFast,
10097
onBuildMax,
@@ -130,7 +127,6 @@ const AgentBody = memo(
130127
nextIndex={nextIndex}
131128
siblingBlocks={p.nestedBlocks}
132129
availableWidth={p.availableWidth}
133-
streamingAgents={p.streamingAgents}
134130
onToggleCollapsed={p.onToggleCollapsed}
135131
markdownPalette={p.markdownPalette}
136132
/>
@@ -157,14 +153,12 @@ const AgentBody = memo(
157153
agentBlocks={agentBlocks}
158154
keyPrefix={`${p.keyPrefix}-agent-grid-${startIndex}`}
159155
availableWidth={p.availableWidth}
160-
streamingAgents={p.streamingAgents}
161156
renderAgentBranch={(innerAgentBlock, prefix, width) => (
162157
<AgentBranchWrapper
163158
agentBlock={innerAgentBlock}
164159
keyPrefix={prefix}
165160
availableWidth={width}
166161
markdownPalette={p.markdownPalette}
167-
streamingAgents={p.streamingAgents}
168162
onToggleCollapsed={p.onToggleCollapsed}
169163
onBuildFast={p.onBuildFast}
170164
onBuildMax={p.onBuildMax}
@@ -249,7 +243,6 @@ export interface AgentBranchWrapperProps {
249243
keyPrefix: string
250244
availableWidth: number
251245
markdownPalette: MarkdownPalette
252-
streamingAgents: Set<string>
253246
onToggleCollapsed: (id: string) => void
254247
onBuildFast: () => void
255248
onBuildMax: () => void
@@ -263,19 +256,18 @@ export const AgentBranchWrapper = memo(
263256
keyPrefix,
264257
availableWidth,
265258
markdownPalette,
266-
streamingAgents,
267259
onToggleCollapsed,
268260
onBuildFast,
269261
onBuildMax,
270262
siblingBlocks,
271263
isLastMessage,
272264
}: AgentBranchWrapperProps) => {
273265
const theme = useTheme()
266+
// Derive streaming boolean for this specific agent to avoid re-renders when other agents change
267+
const agentIsStreaming = useChatStore((state) => state.streamingAgents.has(agentBlock.agentId))
274268

275269
if (shouldRenderAsSimpleText(agentBlock.agentType)) {
276-
const isStreaming =
277-
agentBlock.status === 'running' ||
278-
streamingAgents.has(agentBlock.agentId)
270+
const isStreaming = agentBlock.status === 'running' || agentIsStreaming
279271

280272
const effectiveStatus = isStreaming ? 'running' : agentBlock.status
281273
const { indicator: statusIndicator, color: statusColor } =
@@ -343,8 +335,7 @@ export const AgentBranchWrapper = memo(
343335
}
344336

345337
const isCollapsed = agentBlock.isCollapsed ?? false
346-
const isStreaming =
347-
agentBlock.status === 'running' || streamingAgents.has(agentBlock.agentId)
338+
const isStreaming = agentBlock.status === 'running' || agentIsStreaming
348339

349340
const allTextContent =
350341
agentBlock.blocks
@@ -395,7 +386,6 @@ export const AgentBranchWrapper = memo(
395386
parentIsStreaming={isStreaming}
396387
availableWidth={availableWidth}
397388
markdownPalette={markdownPalette}
398-
streamingAgents={streamingAgents}
399389
onToggleCollapsed={onToggleCollapsed}
400390
onBuildFast={onBuildFast}
401391
onBuildMax={onBuildMax}

‎cli/src/components/blocks/blocks-renderer.tsx‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ interface BlocksRendererProps {
2020
textColor: string
2121
availableWidth: number
2222
markdownPalette: MarkdownPalette
23-
streamingAgents: Set<string>
2423
onToggleCollapsed: (id: string) => void
2524
onBuildFast: () => void
2625
onBuildMax: () => void
@@ -38,7 +37,6 @@ interface BlocksRendererPropsRef {
3837
textColor: string
3938
availableWidth: number
4039
markdownPalette: MarkdownPalette
41-
streamingAgents: Set<string>
4240
onToggleCollapsed: (id: string) => void
4341
onBuildFast: () => void
4442
onBuildMax: () => void
@@ -57,7 +55,6 @@ export const BlocksRenderer = memo(
5755
textColor,
5856
availableWidth,
5957
markdownPalette,
60-
streamingAgents,
6158
onToggleCollapsed,
6259
onBuildFast,
6360
onBuildMax,
@@ -83,7 +80,6 @@ export const BlocksRenderer = memo(
8380
textColor,
8481
availableWidth,
8582
markdownPalette,
86-
streamingAgents,
8783
onToggleCollapsed,
8884
onBuildFast,
8985
onBuildMax,
@@ -130,7 +126,6 @@ export const BlocksRenderer = memo(
130126
nextIndex={nextIndex}
131127
siblingBlocks={p.sourceBlocks}
132128
availableWidth={p.availableWidth}
133-
streamingAgents={p.streamingAgents}
134129
onToggleCollapsed={p.onToggleCollapsed}
135130
markdownPalette={p.markdownPalette}
136131
/>
@@ -157,14 +152,12 @@ export const BlocksRenderer = memo(
157152
agentBlocks={agentBlocks}
158153
keyPrefix={`${p.messageId}-agent-grid-${startIndex}`}
159154
availableWidth={p.availableWidth}
160-
streamingAgents={p.streamingAgents}
161155
renderAgentBranch={(agentBlock, prefix, width) => (
162156
<AgentBranchWrapper
163157
agentBlock={agentBlock}
164158
keyPrefix={prefix}
165159
availableWidth={width}
166160
markdownPalette={p.markdownPalette}
167-
streamingAgents={p.streamingAgents}
168161
onToggleCollapsed={p.onToggleCollapsed}
169162
onBuildFast={p.onBuildFast}
170163
onBuildMax={p.onBuildMax}
@@ -191,7 +184,6 @@ export const BlocksRenderer = memo(
191184
textColor={p.textColor}
192185
availableWidth={p.availableWidth}
193186
markdownPalette={p.markdownPalette}
194-
streamingAgents={p.streamingAgents}
195187
onToggleCollapsed={p.onToggleCollapsed}
196188
onBuildFast={p.onBuildFast}
197189
onBuildMax={p.onBuildMax}

‎cli/src/components/blocks/single-block.tsx‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ interface SingleBlockProps {
2929
textColor: string
3030
availableWidth: number
3131
markdownPalette: MarkdownPalette
32-
streamingAgents: Set<string>
3332
onToggleCollapsed: (id: string) => void
3433
onBuildFast: () => void
3534
onBuildMax: () => void
@@ -49,7 +48,6 @@ export const SingleBlock = memo(
4948
textColor,
5049
availableWidth,
5150
markdownPalette,
52-
streamingAgents,
5351
onToggleCollapsed,
5452
onBuildFast,
5553
onBuildMax,
@@ -176,7 +174,6 @@ export const SingleBlock = memo(
176174
keyPrefix={`${messageId}-agent-${block.agentId}`}
177175
availableWidth={availableWidth}
178176
markdownPalette={markdownPalette}
179-
streamingAgents={streamingAgents}
180177
onToggleCollapsed={onToggleCollapsed}
181178
onBuildFast={onBuildFast}
182179
onBuildMax={onBuildMax}

‎cli/src/components/blocks/tool-block-group.tsx‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ interface ToolBlockGroupProps {
1111
nextIndex: number
1212
siblingBlocks: ContentBlock[]
1313
availableWidth: number
14-
streamingAgents: Set<string>
1514
onToggleCollapsed: (id: string) => void
1615
markdownPalette: MarkdownPalette
1716
}
@@ -50,7 +49,6 @@ export const ToolBlockGroup = memo(
5049
nextIndex,
5150
siblingBlocks,
5251
availableWidth,
53-
streamingAgents,
5452
onToggleCollapsed,
5553
markdownPalette,
5654
}: ToolBlockGroupProps): ReactNode => {
@@ -61,7 +59,6 @@ export const ToolBlockGroup = memo(
6159
toolBlock={toolBlock}
6260
keyPrefix={`${keyPrefix}-tool-${toolBlock.toolCallId}`}
6361
availableWidth={availableWidth}
64-
streamingAgents={streamingAgents}
6562
onToggleCollapsed={onToggleCollapsed}
6663
markdownPalette={markdownPalette}
6764
/>

‎cli/src/components/blocks/tool-branch.tsx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { memo, useCallback } from 'react'
22

33
import { ContentWithMarkdown } from './content-with-markdown'
44
import { useTheme } from '../../hooks/use-theme'
5+
import { useChatStore } from '../../state/chat-store'
56
import { getToolDisplayInfo } from '../../utils/codebuff-client'
67
import { shouldCollapseToolByDefault } from '../../utils/constants'
78
import { renderToolComponent } from '../tools/registry'
@@ -14,7 +15,6 @@ interface ToolBranchProps {
1415
toolBlock: Extract<ContentBlock, { type: 'tool' }>
1516
keyPrefix: string
1617
availableWidth: number
17-
streamingAgents: Set<string>
1818
onToggleCollapsed: (id: string) => void
1919
markdownPalette: MarkdownPalette
2020
}
@@ -24,11 +24,12 @@ export const ToolBranch = memo(
2424
toolBlock,
2525
keyPrefix,
2626
availableWidth,
27-
streamingAgents,
2827
onToggleCollapsed,
2928
markdownPalette,
3029
}: ToolBranchProps) => {
3130
const theme = useTheme()
31+
// Derive streaming boolean for this specific tool to avoid re-renders when other tools/agents change
32+
const isStreaming = useChatStore((state) => state.streamingAgents.has(toolBlock.toolCallId))
3233

3334
const sanitizePreview = (value: string): string =>
3435
value.replace(/[#*_`~\[\]()]/g, '').trim()
@@ -45,7 +46,6 @@ export const ToolBranch = memo(
4546

4647
const displayInfo = getToolDisplayInfo(toolBlock.toolName)
4748
const isCollapsed = toolBlock.isCollapsed ?? shouldCollapseToolByDefault(toolBlock.toolName)
48-
const isStreaming = streamingAgents.has(toolBlock.toolCallId)
4949

5050
const inputContent = `\`\`\`json\n${JSON.stringify(toolBlock.input, null, 2)}\n\`\`\``
5151
const codeBlockLang =

‎cli/src/components/message-block.tsx‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { TextAttributes } from '@opentui/core'
2-
import React, { useState } from 'react'
2+
import { memo, useState } from 'react'
33

44
import { Button } from './button'
55
import { ImageCard } from './image-card'
@@ -40,7 +40,6 @@ interface MessageBlockProps {
4040
markdownOptions: { codeBlockWidth: number; palette: MarkdownPalette }
4141
availableWidth: number
4242
markdownPalette: MarkdownPalette
43-
streamingAgents: Set<string>
4443
onToggleCollapsed: (id: string) => void
4544
onBuildFast: () => void
4645
onBuildMax: () => void
@@ -60,7 +59,7 @@ interface MessageBlockProps {
6059
isLastMessage?: boolean
6160
}
6261

63-
const MessageAttachments = ({
62+
const MessageAttachments = memo(({
6463
imageAttachments,
6564
textAttachments,
6665
}: {
@@ -96,9 +95,9 @@ const MessageAttachments = ({
9695
))}
9796
</box>
9897
)
99-
}
98+
})
10099

101-
export const MessageBlock: React.FC<MessageBlockProps> = ({
100+
export const MessageBlock = memo(({
102101
messageId,
103102
blocks,
104103
content,
@@ -115,7 +114,6 @@ export const MessageBlock: React.FC<MessageBlockProps> = ({
115114
markdownOptions,
116115
availableWidth,
117116
markdownPalette,
118-
streamingAgents,
119117
onToggleCollapsed,
120118
onBuildFast,
121119
onBuildMax,
@@ -128,7 +126,7 @@ export const MessageBlock: React.FC<MessageBlockProps> = ({
128126
textAttachments,
129127
metadata,
130128
isLastMessage,
131-
}) => {
129+
}: MessageBlockProps) => {
132130
const [showValidationPopover, setShowValidationPopover] = useState(false)
133131

134132
const bashCwd = metadata?.bashCwd ? formatCwd(metadata.bashCwd) : undefined
@@ -153,7 +151,6 @@ export const MessageBlock: React.FC<MessageBlockProps> = ({
153151
markdownOptions,
154152
availableWidth,
155153
markdownPalette,
156-
streamingAgents,
157154
onToggleCollapsed,
158155
onBuildFast,
159156
onBuildMax,
@@ -276,7 +273,6 @@ export const MessageBlock: React.FC<MessageBlockProps> = ({
276273
textColor={resolvedTextColor}
277274
availableWidth={availableWidth}
278275
markdownPalette={markdownPalette}
279-
streamingAgents={streamingAgents}
280276
onToggleCollapsed={onToggleCollapsed}
281277
onBuildFast={onBuildFast}
282278
onBuildMax={onBuildMax}
@@ -326,4 +322,4 @@ export const MessageBlock: React.FC<MessageBlockProps> = ({
326322
)}
327323
</box>
328324
)
329-
}
325+
})

0 commit comments

Comments
 (0)