Skip to content

Commit 0274a48

Browse files
brandonkachenbrandonkachen
committed
Style CLI tool branches like tree
1 parent 660736b commit 0274a48

File tree

2 files changed

+138
-49
lines changed

2 files changed

+138
-49
lines changed

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

Lines changed: 46 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import React, { type ReactNode } from 'react'
44
import { pluralize } from '@codebuff/common/util/string'
55

66
import { BranchItem } from './branch-item'
7-
import { ToolItem } from './tool-item'
7+
import { ToolItem, type ToolBranchMeta } from './tool-item'
88
import { getToolRenderConfig } from './tool-renderer'
99
import { getToolDisplayInfo } from '../utils/codebuff-client'
1010
import {
@@ -81,10 +81,33 @@ export const MessageBlock = ({
8181
}
8282
}
8383

84+
const defaultToolBranchMeta: ToolBranchMeta = {
85+
hasPrevious: false,
86+
hasNext: false,
87+
}
88+
89+
const buildToolBranchMeta = (
90+
items: Array<ContentBlock | undefined>,
91+
): Map<number, ToolBranchMeta> => {
92+
const toolIndices = items
93+
.map((block, index) => (block && block.type === 'tool' ? index : -1))
94+
.filter((index) => index !== -1) as number[]
95+
96+
const meta = new Map<number, ToolBranchMeta>()
97+
toolIndices.forEach((blockIndex, position) => {
98+
meta.set(blockIndex, {
99+
hasPrevious: position > 0,
100+
hasNext: position < toolIndices.length - 1,
101+
})
102+
})
103+
return meta
104+
}
105+
84106
const renderToolBranch = (
85107
toolBlock: Extract<ContentBlock, { type: 'tool' }>,
86108
indentLevel: number,
87109
keyPrefix: string,
110+
branchMeta: ToolBranchMeta,
88111
): React.ReactNode => {
89112
if (toolBlock.toolName === 'end_turn') {
90113
return null
@@ -170,6 +193,7 @@ export const MessageBlock = ({
170193
streamingPreview={streamingPreview}
171194
finishedPreview={finishedPreview}
172195
theme={theme}
196+
branchMeta={branchMeta}
173197
onToggle={() => onToggleCollapsed(toolBlock.toolCallId)}
174198
/>
175199
</box>
@@ -333,6 +357,7 @@ export const MessageBlock = ({
333357
): React.ReactNode[] {
334358
const nestedBlocks = agentBlock.blocks ?? []
335359
const nodes: React.ReactNode[] = []
360+
const toolBranchMetaMap = buildToolBranchMeta(nestedBlocks)
336361

337362
nestedBlocks.forEach((nestedBlock, nestedIdx) => {
338363
if (nestedBlock.type === 'text') {
@@ -364,11 +389,14 @@ export const MessageBlock = ({
364389
</text>,
365390
)
366391
} else if (nestedBlock.type === 'tool') {
392+
const branchMeta =
393+
toolBranchMetaMap.get(nestedIdx) ?? defaultToolBranchMeta
367394
nodes.push(
368395
renderToolBranch(
369396
nestedBlock,
370397
indentLevel,
371398
`${keyPrefix}-tool-${nestedBlock.toolCallId}`,
399+
branchMeta,
372400
),
373401
)
374402
} else if (nestedBlock.type === 'agent') {
@@ -385,6 +413,8 @@ export const MessageBlock = ({
385413
return nodes
386414
}
387415

416+
const topLevelToolMeta = blocks ? buildToolBranchMeta(blocks) : null
417+
388418
return (
389419
<>
390420
{isUser && (
@@ -422,18 +452,21 @@ export const MessageBlock = ({
422452
? 0
423453
: 0
424454
const blockTextColor = block.color ?? textColor
425-
return (
426-
<text key={renderKey} style={{ fg: blockTextColor, marginTop }}>
427-
{renderedContent}
428-
</text>
429-
)
430-
} else if (block.type === 'tool') {
431-
return renderToolBranch(
432-
block,
433-
0,
434-
`${messageId}-tool-${block.toolCallId}`,
435-
)
436-
} else if (block.type === 'agent') {
455+
return (
456+
<text key={renderKey} style={{ fg: blockTextColor, marginTop }}>
457+
{renderedContent}
458+
</text>
459+
)
460+
} else if (block.type === 'tool') {
461+
const branchMeta =
462+
topLevelToolMeta?.get(idx) ?? defaultToolBranchMeta
463+
return renderToolBranch(
464+
block,
465+
0,
466+
`${messageId}-tool-${block.toolCallId}`,
467+
branchMeta,
468+
)
469+
} else if (block.type === 'agent') {
437470
return renderAgentBranch(
438471
block,
439472
0,

‎cli/src/components/tool-item.tsx‎

Lines changed: 92 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,11 @@ import React, { type ReactNode } from 'react'
33

44
import type { ChatTheme } from '../utils/theme-system'
55

6+
export interface ToolBranchMeta {
7+
hasPrevious: boolean
8+
hasNext: boolean
9+
}
10+
611
interface ToolItemProps {
712
name: string
813
titleAccessory?: ReactNode
@@ -12,6 +17,7 @@ interface ToolItemProps {
1217
streamingPreview: string
1318
finishedPreview: string
1419
theme: ChatTheme
20+
branchMeta: ToolBranchMeta
1521
onToggle: () => void
1622
}
1723

@@ -65,18 +71,93 @@ export const ToolItem = ({
6571
streamingPreview,
6672
finishedPreview,
6773
theme,
74+
branchMeta,
6875
onToggle,
6976
}: ToolItemProps) => {
70-
const toggleColor = theme.statusSecondary
71-
const toggleIcon = isCollapsed ? '▸' : '▾'
77+
const branchColor = theme.agentResponseCount
78+
const branchAttributes = TextAttributes.DIM
79+
const titleColor = theme.statusSecondary
7280
const previewColor = isStreaming ? theme.agentText : theme.agentResponseCount
81+
const connectorSymbol = branchMeta.hasNext ? '├' : '└'
82+
const continuationPrefix = branchMeta.hasNext ? '│ ' : ' '
83+
const showBranchAbove = branchMeta.hasPrevious
7384
const hasTitleAccessory =
74-
titleAccessory !== undefined &&
75-
titleAccessory !== null &&
76-
!(typeof titleAccessory === 'string' && titleAccessory.length === 0)
85+
titleAccessory !== undefined && titleAccessory !== null
86+
87+
const renderBranchSpacer = () => {
88+
if (!showBranchAbove) {
89+
return null
90+
}
91+
92+
return (
93+
<box
94+
style={{
95+
flexDirection: 'row',
96+
paddingLeft: 1,
97+
paddingRight: 1,
98+
paddingTop: 0,
99+
paddingBottom: 0,
100+
}}
101+
>
102+
<text style={{ wrapMode: 'none' }}>
103+
<span fg={branchColor} attributes={branchAttributes}>
104+
│
105+
</span>
106+
</text>
107+
</box>
108+
)
109+
}
110+
111+
const renderConnectedSection = (node: ReactNode) => {
112+
if (!node) {
113+
return null
114+
}
115+
116+
return (
117+
<box
118+
style={{
119+
flexDirection: 'row',
120+
gap: 0,
121+
paddingLeft: 1,
122+
paddingRight: 1,
123+
paddingTop: 0,
124+
paddingBottom: 0,
125+
}}
126+
>
127+
<text style={{ wrapMode: 'none' }}>
128+
<span fg={branchColor} attributes={branchAttributes}>
129+
{continuationPrefix}
130+
</span>
131+
</text>
132+
<box
133+
style={{
134+
flexDirection: 'column',
135+
gap: 0,
136+
paddingLeft: 0,
137+
paddingRight: 0,
138+
paddingTop: 0,
139+
paddingBottom: 0,
140+
}}
141+
>
142+
{node}
143+
</box>
144+
</box>
145+
)
146+
}
147+
148+
const renderedContent = renderContent(content, theme)
149+
const previewText = isStreaming ? streamingPreview : finishedPreview
150+
const hasPreview =
151+
typeof previewText === 'string' ? previewText.length > 0 : false
152+
const previewNode = hasPreview ? (
153+
<text fg={previewColor} attributes={TextAttributes.ITALIC}>
154+
{previewText}
155+
</text>
156+
) : null
77157

78158
return (
79159
<box style={{ flexDirection: 'column', gap: 0 }}>
160+
{renderBranchSpacer()}
80161
<box
81162
style={{
82163
flexDirection: 'row',
@@ -89,8 +170,10 @@ export const ToolItem = ({
89170
onMouseDown={onToggle}
90171
>
91172
<text style={{ wrapMode: 'none' }}>
92-
<span fg={toggleColor}>{toggleIcon} </span>
93-
<span fg={toggleColor} attributes={TextAttributes.BOLD}>
173+
<span fg={branchColor} attributes={branchAttributes}>
174+
{connectorSymbol}{' '}
175+
</span>
176+
<span fg={titleColor} attributes={TextAttributes.BOLD}>
94177
{name}
95178
</span>
96179
{hasTitleAccessory ? (
@@ -101,35 +184,8 @@ export const ToolItem = ({
101184
) : null}
102185
</text>
103186
</box>
104-
{isCollapsed ? (
105-
(isStreaming && streamingPreview) || (!isStreaming && finishedPreview) ? (
106-
<box
107-
style={{
108-
paddingLeft: 3,
109-
paddingRight: 1,
110-
paddingTop: 0,
111-
paddingBottom: 0,
112-
}}
113-
>
114-
<text fg={previewColor} attributes={TextAttributes.ITALIC}>
115-
{isStreaming ? streamingPreview : finishedPreview}
116-
</text>
117-
</box>
118-
) : null
119-
) : (
120-
<box
121-
style={{
122-
flexDirection: 'column',
123-
gap: 0,
124-
paddingLeft: 3,
125-
paddingRight: 1,
126-
paddingTop: 0,
127-
paddingBottom: 0,
128-
}}
129-
>
130-
{renderContent(content, theme)}
131-
</box>
132-
)}
187+
{isCollapsed ? renderConnectedSection(previewNode) : null}
188+
{!isCollapsed ? renderConnectedSection(renderedContent) : null}
133189
</box>
134190
)
135191
}

0 commit comments

Comments
 (0)