From 4bdd0d5ea7e6b5f9be211de40280083aea5e8ae1 Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian <157918347+Xiaohan-Tian@users.noreply.github.com> Date: Tue, 2 Jun 2026 22:31:29 -0700 Subject: [PATCH] feat: minor UI adjustments --- src/components/ChatBox.css | 87 +++++++++ src/components/ChatBox.tsx | 3 + src/components/chat/AssistantMessage.test.tsx | 97 ++++++++++ src/components/chat/AssistantMessage.tsx | 137 ++++++++++++++- src/hooks/useStreamProcessor.test.ts | 108 ++++++++++++ src/hooks/useStreamProcessor.ts | 165 +++++++++++++++++- src/types/projectTypes.ts | 3 + 7 files changed, 593 insertions(+), 7 deletions(-) diff --git a/src/components/ChatBox.css b/src/components/ChatBox.css index 65a3f3f..7169ce4 100644 --- a/src/components/ChatBox.css +++ b/src/components/ChatBox.css @@ -398,12 +398,99 @@ font-weight: bold; } +.message-tool-result-title, +.message-tool-summary { + margin: 4px 0; +} + +.message-tool-summary { + display: flex; + align-items: flex-start; + gap: 0; +} + +.message-tool-summary-prefix { + white-space: pre; + flex: 0 0 auto; +} + +.message-tool-summary-content { + flex: 1 1 auto; +} + +.message-tool-summary-content > :first-child { + margin-top: 0; +} + +.message-tool-summary-content > :last-child { + margin-bottom: 0; +} + .message-performance-info { margin-top: 8px; font-size: 10px; color: #909090; } +.tool-call-code-block { + position: relative; + margin: 8px 0; +} + +.tool-call-code-block-inner { + overflow: hidden; + transition: max-height 0.24s ease; +} + +.tool-call-code-block-inner > div { + margin: 0 !important; +} + +.tool-call-code-block-inner pre { + margin: 0 !important; + overflow-x: auto !important; + overflow-y: hidden !important; +} + +.tool-call-code-block-toggle { + position: absolute; + left: 0; + right: 0; + bottom: 0; + display: flex; + align-items: flex-end; + justify-content: center; + width: 100%; + padding: 0; + border: 0; + background: transparent; + cursor: pointer; +} + +.tool-call-code-block-toggle-content { + position: relative; + z-index: 1; + display: inline-flex; + align-items: center; + gap: 8px; + padding: 20px 12px 10px; + color: #d7e7f6; + font-size: 11px; + font-weight: 600; +} + +.tool-call-code-block-toggle:hover .tool-call-code-block-toggle-content { + color: #f0f7ff; +} + +.tool-call-code-block-fade { + position: absolute; + left: 0; + right: 0; + bottom: 0; + background: linear-gradient(180deg, rgba(26, 26, 26, 0) 0%, rgba(26, 26, 26, 0.92) 58%, #1a1a1a 100%); +} + .message-divider-banner { display: flex; align-items: center; diff --git a/src/components/ChatBox.tsx b/src/components/ChatBox.tsx index e849c6a..f4b1e73 100644 --- a/src/components/ChatBox.tsx +++ b/src/components/ChatBox.tsx @@ -527,7 +527,10 @@ const ChatBox: React.FC = ({ isVisible }) => { performanceInfo={message.performanceInfo} toolName={message.toolName} toolSuccess={message.toolSuccess} + toolRawResult={message.toolRawResult} + toolResultDisplayContent={message.toolResultDisplayContent} todoSnapshot={message.todoSnapshot} + isToolCallMessage={message.isToolCallMessage} onAbort={message.isStreaming ? handleAbort : undefined} /> ) diff --git a/src/components/chat/AssistantMessage.test.tsx b/src/components/chat/AssistantMessage.test.tsx index 04a8e14..f8d73db 100644 --- a/src/components/chat/AssistantMessage.test.tsx +++ b/src/components/chat/AssistantMessage.test.tsx @@ -6,6 +6,7 @@ import type { TodoItem } from '../../agent/core/todo'; describe('AssistantMessage', () => { afterEach(() => { vi.useRealTimers(); + vi.restoreAllMocks(); }); it.each([ @@ -95,6 +96,68 @@ describe('AssistantMessage', () => { expect(codeElement).toHaveTextContent('const value = 1;'); }); + it('keeps non-tool-call code blocks rendered without the expander', () => { + const { container } = render( + + ); + + expect(container.querySelector('[data-testid="tool-call-code-block"]')).not.toBeInTheDocument(); + expect(screen.queryByRole('button', { name: 'Click to expand' })).not.toBeInTheDocument(); + }); + + it('collapses long tool-call code blocks and expands them in place', () => { + vi.spyOn(HTMLElement.prototype, 'scrollHeight', 'get').mockReturnValue(280); + vi.stubGlobal('requestAnimationFrame', (callback: FrameRequestCallback) => { + callback(0); + return 0; + }); + + const { container } = render( +
+ +
+ ); + + const toolCallBlock = screen.getByTestId('tool-call-code-block'); + const toolCallInner = container.querySelector('.tool-call-code-block-inner') as HTMLDivElement; + + expect(toolCallBlock).toBeInTheDocument(); + expect(toolCallInner.style.maxHeight).toBe('200px'); + expect(screen.getByRole('button', { name: 'Click to expand' })).toBeInTheDocument(); + + fireEvent.click(screen.getByRole('button', { name: 'Click to expand' })); + + expect(screen.getByRole('button', { name: 'Click to collapse' })).toBeInTheDocument(); + expect(toolCallInner.style.maxHeight).toBe('280px'); + }); + + it('preserves chat scroll position when expanding a tool-call block', () => { + vi.spyOn(HTMLElement.prototype, 'scrollHeight', 'get').mockReturnValue(320); + vi.stubGlobal('requestAnimationFrame', (callback: FrameRequestCallback) => { + callback(0); + return 0; + }); + + render( +
+ +
+ ); + + const scrollContainer = document.querySelector('.chatbox-messages') as HTMLDivElement; + scrollContainer.scrollTop = 96; + + fireEvent.click(screen.getByRole('button', { name: 'Click to expand' })); + + expect(scrollContainer.scrollTop).toBe(96); + }); + it('renders compacting and compacted messages as divider banners', () => { const { rerender, container } = render( @@ -133,4 +196,38 @@ describe('AssistantMessage', () => { expect(screen.getByText('Working on: Writing harmony')).toBeInTheDocument(); expect(screen.queryByText('fallback content')).not.toBeInTheDocument(); }); + + it('renders the add_notes summary instead of the raw tool result text', () => { + render( + + ); + + expect(screen.getByText('add_notes')).toBeInTheDocument(); + expect(screen.getByText(/└──/)).toBeInTheDocument(); + expect(screen.getByText(/Successfully created 81 notes in region/i)).toBeInTheDocument(); + expect(screen.getByText('Verse Melody')).toBeInTheDocument(); + expect(screen.getByText('Lead Vox')).toBeInTheDocument(); + expect(screen.queryByText(/Successfully created 81 notes: C4/)).not.toBeInTheDocument(); + }); + + it('renders generic tool results with the stable shell and raw body by default', () => { + render( + + ); + + expect(screen.getByText('read_music')).toBeInTheDocument(); + expect(screen.getByText(/└──/)).toBeInTheDocument(); + expect(screen.getByText('C D E F')).toBeInTheDocument(); + }); }); diff --git a/src/components/chat/AssistantMessage.tsx b/src/components/chat/AssistantMessage.tsx index 7a8e532..5b705a6 100644 --- a/src/components/chat/AssistantMessage.tsx +++ b/src/components/chat/AssistantMessage.tsx @@ -1,8 +1,9 @@ -import React, { memo, useEffect, useState } from 'react'; +import React, { memo, useEffect, useRef, useState } from 'react'; import ReactMarkdown from 'react-markdown'; import rehypeKatex from 'rehype-katex'; import remarkGfm from 'remark-gfm'; import remarkMath from 'remark-math'; +import { FaCaretDown, FaCaretUp } from 'react-icons/fa'; import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'; import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism'; import type { PerformanceInfo } from '../../agent/llm/StreamingTypes'; @@ -16,13 +17,109 @@ interface AssistantMessageProps { performanceInfo?: PerformanceInfo; toolName?: string; toolSuccess?: boolean; + toolRawResult?: string; + toolResultDisplayContent?: string; todoSnapshot?: TodoItem[]; + isToolCallMessage?: boolean; } -// Memoized code component to prevent SyntaxHighlighter re-renders -// eslint-disable-next-line @typescript-eslint/no-explicit-any -const CodeComponent = memo(({ inline, className, children, ...props }: any) => { +const COLLAPSED_TOOL_CALL_HEIGHT_PX = 200; +const TOOL_CALL_FADE_HEIGHT_PX = 50; + +interface MarkdownCodeProps { + inline?: boolean; + className?: string; + children?: React.ReactNode; +} + +const ToolCallCodeBlock = memo(({ + className, + children, +}: Omit) => { + const containerRef = useRef(null); + const [isExpandable, setIsExpandable] = useState(false); + const [isExpanded, setIsExpanded] = useState(false); + const [expandedHeight, setExpandedHeight] = useState(COLLAPSED_TOOL_CALL_HEIGHT_PX); const match = /language-(\w+)/.exec(className || ''); + const codeText = String(children).replace(/\n$/, ''); + + useEffect(() => { + const container = containerRef.current; + if (!container) { + return; + } + + const measuredHeight = container.scrollHeight; + setExpandedHeight(measuredHeight); + setIsExpandable(measuredHeight > COLLAPSED_TOOL_CALL_HEIGHT_PX); + }, [codeText]); + + const handleToggleExpanded = () => { + const scrollContainer = containerRef.current?.closest('.chatbox-messages') as HTMLDivElement | null; + const previousScrollTop = scrollContainer?.scrollTop ?? null; + + setIsExpanded((current) => !current); + + window.requestAnimationFrame(() => { + if (scrollContainer && previousScrollTop !== null) { + scrollContainer.scrollTop = previousScrollTop; + } + }); + }; + + const maxHeight = isExpanded ? `${expandedHeight}px` : `${COLLAPSED_TOOL_CALL_HEIGHT_PX}px`; + + return ( +
+
+ + {codeText} + +
+ {isExpandable && ( + + )} +
+ ); +}); + +// Memoized code component to prevent SyntaxHighlighter re-renders +const CodeComponent = memo(({ inline, className, children, isToolCallMessage, ...props }: MarkdownCodeProps & { isToolCallMessage: boolean }) => { + const match = /language-(\w+)/.exec(className || ''); + if (!inline && match && isToolCallMessage) { + return ( + + {children} + + ); + } + return !inline && match ? ( = ({ performanceInfo, toolName, toolSuccess, + toolRawResult, + toolResultDisplayContent, todoSnapshot, + isToolCallMessage = false, }) => { const prefillTps = formatTps(performanceInfo?.prefillTps); const generationTps = formatTps(performanceInfo?.generationTps); @@ -82,6 +182,8 @@ const AssistantMessage: React.FC = ({ || content === COMPACTION_DONE_LABEL || content === COMPACTION_EMPTY_LABEL; const isTodoSnapshotCard = toolName === 'update_todo_list' && Array.isArray(todoSnapshot); + const shouldRenderGenericToolResult = Boolean(toolName) && typeof toolSuccess === 'boolean' && !isTodoSnapshotCard; + const genericToolDisplayContent = toolResultDisplayContent ?? toolRawResult ?? content; useEffect(() => { if (!isThinking) { @@ -184,12 +286,37 @@ const AssistantMessage: React.FC = ({ } } + if (shouldRenderGenericToolResult && toolName) { + return ( +
+

+ {' '} + {toolName} +

+
+ +
+ , + }} + > + {genericToolDisplayContent} + +
+
+
+ ); + } + return ( , }} > {content} diff --git a/src/hooks/useStreamProcessor.test.ts b/src/hooks/useStreamProcessor.test.ts index 5e6c8ba..97f6eb3 100644 --- a/src/hooks/useStreamProcessor.test.ts +++ b/src/hooks/useStreamProcessor.test.ts @@ -8,6 +8,24 @@ vi.mock('../agent/core/AgentCore', () => ({ } })); +vi.mock('../core/KGCore', () => ({ + KGCore: { + instance: vi.fn() + } +})); + +vi.mock('../stores/projectStore', () => ({ + useProjectStore: { + getState: vi.fn(() => ({ + activeRegionId: null, + selectedRegionIds: [], + timeSignature: { numerator: 4, denominator: 4 }, + tracks: [], + refreshProjectState: vi.fn(), + })), + }, +})); + vi.mock('../utils/chatMessageUtils', () => ({ createStreamingMessage: () => ({ id: 'streaming-message', @@ -24,6 +42,8 @@ vi.mock('../utils/chatMessageUtils', () => ({ })); import { AgentCore } from '../agent/core/AgentCore'; +import { KGCore } from '../core/KGCore'; +import { KGMidiRegion } from '../core/region/KGMidiRegion'; import { useStreamProcessor } from './useStreamProcessor'; import type { ChatMessage } from '../types/projectTypes'; @@ -215,6 +235,94 @@ describe('useStreamProcessor', () => { const addedMessages = [...messages.values()]; expect(addedMessages.some(message => message.content.includes('Calling tool: read_music'))).toBe(true); + expect(addedMessages.some(message => message.isToolCallMessage)).toBe(true); expect(addedMessages.some(message => message.toolName === 'update_todo_list')).toBe(false); }); + +it('attaches add_notes summary metadata for chat-only rendering while preserving raw content', async () => { + const selectedRegion = new KGMidiRegion('region-1', '1', 0, 'Verse Melody'); + + vi.spyOn(AgentCore, 'instance').mockReturnValue({ + getAgentState: () => ({ + getTodos: () => [], + }), + processUserInput: async function* () { + yield { + type: 'tool_call', + content: '', + toolCall: { + id: 'add-notes-call-1', + type: 'function', + function: { + name: 'add_notes', + arguments: JSON.stringify({ + notes: [ + { pitch: 'C4', start: 16, length: 4 }, + { pitch: 'E4', start: 20, length: 8 }, + ], + }), + }, + }, + }; + yield { + type: 'tool_result', + content: '', + toolResult: { + name: 'add_notes', + success: true, + result: 'Successfully created 2 notes: C4 (beat 16, length 4), E4 (beat 20, length 8)', + }, + }; + yield { type: 'done', content: '' }; + }, + } as unknown as AgentCore); + + vi.spyOn(KGCore, 'instance').mockReturnValue({ + getCurrentProject: () => ({ + getTimeSignature: () => ({ numerator: 4, denominator: 4 }), + getTracks: () => [ + { + getId: () => '1', + getName: () => 'Lead', + getRegions: () => [selectedRegion], + }, + ], + }), + getSelectedItems: () => [selectedRegion], + } as unknown as KGCore); + + const messages = new Map(); + + const { result } = renderHook(() => useStreamProcessor({ + onMessageAdd: (message) => { + messages.set(message.id, message); + }, + onMessageUpdate: (messageId, updater) => { + const current = messages.get(messageId); + if (!current) { + throw new Error(`Missing message ${messageId}`); + } + messages.set(messageId, updater(current)); + }, + onMessageRemove: (messageId) => { + messages.delete(messageId); + }, + onProcessingChange: () => undefined, + })); + + await act(async () => { + await result.current.processStream('add notes prompt'); + }); + + const addedMessages = [...messages.values()]; + const toolCallMessage = addedMessages.find(message => message.isToolCallMessage); + const addNotesMessage = addedMessages.find(message => message.toolName === 'add_notes'); + + expect(toolCallMessage?.content).toContain('Calling tool: add_notes'); + expect(addNotesMessage?.content).toContain('Successfully created 2 notes:'); + expect(addNotesMessage?.toolRawResult).toBe('Successfully created 2 notes: C4 (beat 16, length 4), E4 (beat 20, length 8)'); + expect(addNotesMessage?.toolResultDisplayContent).toBe( + 'Successfully created 2 notes in region **Verse Melody** on track **Lead**, spanning bars 5 to 7.' + ); + }); }); diff --git a/src/hooks/useStreamProcessor.ts b/src/hooks/useStreamProcessor.ts index 285afe4..0f55c74 100644 --- a/src/hooks/useStreamProcessor.ts +++ b/src/hooks/useStreamProcessor.ts @@ -1,9 +1,148 @@ import { useState, useCallback } from 'react'; import { AgentCore } from '../agent/core/AgentCore'; +import { KGCore } from '../core/KGCore'; +import { KGMidiRegion } from '../core/region/KGMidiRegion'; import { createStreamingMessage, createMessage } from '../utils/chatMessageUtils'; +import { useProjectStore } from '../stores/projectStore'; import type { ChatMessage } from '../types/projectTypes'; const TODO_TOOL_NAME = 'update_todo_list'; +const ADD_NOTES_TOOL_NAME = 'add_notes'; + +interface PendingToolCall { + name: string; + arguments: Record | null; +} + +interface AddNotesToolArguments { + notes: Array<{ + pitch: string; + start: number; + length: number; + velocity?: number; + }>; + region_id?: string; +} + +const getBarNumberFromStartBeat = (beat: number, beatsPerBar: number): number => ( + Math.floor(beat / beatsPerBar) + 1 +); + +const getBarNumberFromEndBeat = (beat: number, beatsPerBar: number): number => ( + Math.max(1, Math.ceil(beat / beatsPerBar)) +); + +interface AddNotesSummaryData { + noteCount: number; + regionName: string; + trackName: string; + earliestNoteStartBar: number; + latestNoteEndBar: number; +} + +const resolveTargetMidiRegion = ( + regionId: string | undefined, + storeState: ReturnType, +): { region: KGMidiRegion; trackName: string } | undefined => { + const project = KGCore.instance().getCurrentProject(); + const tracks = project.getTracks(); + + const findById = (candidateRegionId: string): { region: KGMidiRegion; trackName: string } | undefined => { + for (const track of tracks) { + const region = track.getRegions().find(candidate => candidate.getId() === candidateRegionId); + if (region instanceof KGMidiRegion) { + return { + region, + trackName: track.getName(), + }; + } + } + + return undefined; + }; + + if (regionId) { + return findById(regionId); + } + + if (storeState.activeRegionId) { + const activeRegion = findById(storeState.activeRegionId); + if (activeRegion) { + return activeRegion; + } + } + + const selectedRegionId = storeState.selectedRegionIds.at(-1); + if (selectedRegionId) { + const selectedRegion = findById(selectedRegionId); + if (selectedRegion) { + return selectedRegion; + } + } + + const selectedItems = KGCore.instance().getSelectedItems(); + for (const item of selectedItems) { + if (item instanceof KGMidiRegion) { + const track = tracks.find(candidate => candidate.getId() === item.getTrackId()); + return { + region: item, + trackName: track?.getName() ?? `Track ${item.getTrackIndex() + 1}`, + }; + } + } + + return undefined; +}; + +const buildAddNotesSummary = (args: Record | null): AddNotesSummaryData | undefined => { + if (!args) { + return undefined; + } + + const typedArgs = args as AddNotesToolArguments; + if (!Array.isArray(typedArgs.notes) || typedArgs.notes.length === 0) { + return undefined; + } + + const storeState = useProjectStore.getState(); + const project = KGCore.instance().getCurrentProject(); + const beatsPerBar = project.getTimeSignature().numerator ?? storeState.timeSignature.numerator; + const targetRegion = resolveTargetMidiRegion(typedArgs.region_id, storeState); + if (!targetRegion) { + return undefined; + } + + const earliestNoteStartBeat = Math.min(...typedArgs.notes.map(note => note.start)); + const latestNoteEndBeat = Math.max(...typedArgs.notes.map(note => note.start + note.length)); + + return { + noteCount: typedArgs.notes.length, + regionName: targetRegion.region.getName(), + trackName: targetRegion.trackName, + earliestNoteStartBar: getBarNumberFromStartBeat(earliestNoteStartBeat, beatsPerBar), + latestNoteEndBar: getBarNumberFromEndBeat(latestNoteEndBeat, beatsPerBar), + }; +}; + +const buildToolResultDisplayContent = ( + toolName: string, + success: boolean, + rawResult: string, + toolArgs: Record | null, +): string => { + if (!success) { + return rawResult; + } + + if (toolName === ADD_NOTES_TOOL_NAME) { + const summary = buildAddNotesSummary(toolArgs); + if (summary) { + return `Successfully created ${summary.noteCount} ${summary.noteCount === 1 ? 'note' : 'notes'} in region **${summary.regionName}** on track **${summary.trackName}**, spanning bars ${summary.earliestNoteStartBar} to ${summary.latestNoteEndBar}.`; + } + } + + return rawResult; +}; interface StreamProcessorOptions { onMessageUpdate: (messageId: string, updater: (msg: ChatMessage) => ChatMessage) => void; @@ -35,6 +174,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce // Track the current streaming message ID (mutable) const initialStreamingMsg = createStreamingMessage(); let currentStreamingId = initialStreamingMsg.id; + const pendingToolCalls: PendingToolCall[] = []; onMessageAdd(initialStreamingMsg); try { @@ -91,14 +231,29 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce try { const args = JSON.parse(chunk.toolCall.function.arguments); argsDisplay = JSON.stringify(args, null, 2); + pendingToolCalls.push({ name: toolName, arguments: args }); } catch { argsDisplay = chunk.toolCall.function.arguments; + pendingToolCalls.push({ name: toolName, arguments: null }); } - const toolCallMsg = createMessage('assistant', `🔧 **Calling tool: ${toolName}**\n\n\`\`\`json\n${argsDisplay}\n\`\`\``); + const toolCallMsg = { + ...createMessage('assistant', `🔧 **Calling tool: ${toolName}**\n\n\`\`\`json\n${argsDisplay}\n\`\`\``), + isToolCallMessage: true, + }; onMessageAdd(toolCallMsg); } else if (chunk.type === 'tool_result' && chunk.toolResult) { // Show tool result in UI const { name, success, result } = chunk.toolResult; + const pendingToolCallIndex = pendingToolCalls.findIndex(toolCall => toolCall.name === name); + const pendingToolCall = pendingToolCallIndex >= 0 + ? pendingToolCalls.splice(pendingToolCallIndex, 1)[0] + : undefined; + const toolResultDisplayContent = buildToolResultDisplayContent( + name, + success, + result, + pendingToolCall?.arguments ?? null, + ); const toolResultMsg = name === TODO_TOOL_NAME ? { ...createMessage('assistant', result), @@ -106,7 +261,13 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce toolSuccess: success, todoSnapshot: AgentCore.instance().getAgentState().getTodos().map(todo => ({ ...todo })), } - : createMessage('assistant', `${success ? '✅' : '❌'} **${name}**\n\n └── ${result}`); + : { + ...createMessage('assistant', `${success ? '✅' : '❌'} **${name}**\n\n └── ${result}`), + toolName: name, + toolSuccess: success, + toolRawResult: result, + toolResultDisplayContent, + }; onMessageAdd(toolResultMsg); // Reset for the next LLM turn in the agentic loop diff --git a/src/types/projectTypes.ts b/src/types/projectTypes.ts index aaeb91f..6077a60 100644 --- a/src/types/projectTypes.ts +++ b/src/types/projectTypes.ts @@ -16,7 +16,10 @@ export interface ChatMessage { performanceInfo?: PerformanceInfo; toolName?: string; toolSuccess?: boolean; + toolRawResult?: string; + toolResultDisplayContent?: string; todoSnapshot?: TodoItem[]; + isToolCallMessage?: boolean; } /**