fix: switch to processing status when receiving response tokens from LLM

This commit is contained in:
Xiaohan-Tian
2026-05-14 21:37:06 -07:00
parent 1a051fea54
commit 2bd6925f27
5 changed files with 139 additions and 8 deletions
@@ -0,0 +1,25 @@
import { fireEvent, render, screen } from '@testing-library/react';
import { describe, expect, it, vi } from 'vitest';
import AssistantMessage from './AssistantMessage';
describe('AssistantMessage', () => {
it.each([
'<span class="processing-wave">Thinking...</span> click here to abort.',
'<span class="processing-wave">Processing...</span> 3 tokens received. click here to abort.'
])('renders the abort action for streaming status content: %s', (content) => {
const onAbort = vi.fn();
render(
<AssistantMessage
content={content}
isStreaming
onAbort={onAbort}
/>
);
const abortButton = screen.getByRole('button', { name: 'click here to abort' });
expect(abortButton).toBeInTheDocument();
fireEvent.click(abortButton);
expect(onAbort).toHaveBeenCalledTimes(1);
});
});
+8 -4
View File
@@ -44,22 +44,26 @@ const AssistantMessage: React.FC<AssistantMessageProps> = ({ content, isStreamin
const prefillTps = formatTps(performanceInfo?.prefillTps); const prefillTps = formatTps(performanceInfo?.prefillTps);
const generationTps = formatTps(performanceInfo?.generationTps); const generationTps = formatTps(performanceInfo?.generationTps);
const hasPerformanceInfo = Boolean(prefillTps || generationTps); const hasPerformanceInfo = Boolean(prefillTps || generationTps);
const processingWaveLabels = ['Thinking...', 'Processing...'];
const renderContent = () => { const renderContent = () => {
// Handle special abort link for streaming messages // Handle special abort link for streaming messages
if (isStreaming && onAbort && content.includes('click here to abort')) { if (isStreaming && onAbort && content.includes('click here to abort')) {
const hasProcessingWave = content.includes('<span class="processing-wave">Thinking...</span>'); const processingWaveMarkup = processingWaveLabels
.map(label => `<span class="processing-wave">${label}</span>`)
.find(markup => content.includes(markup));
if (hasProcessingWave) { if (processingWaveMarkup) {
const parts = content.split('click here to abort'); const parts = content.split('click here to abort');
const beforeAbort = parts[0].replace( const beforeAbort = parts[0].replace(
'<span class="processing-wave">Thinking...</span>', processingWaveMarkup,
'' ''
); );
const waveLabel = processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? 'Thinking...';
return ( return (
<span> <span>
<span className="processing-wave">Thinking...</span> <span className="processing-wave">{waveLabel}</span>
{beforeAbort} {beforeAbort}
<button onClick={onAbort} className="abort-link"> <button onClick={onAbort} className="abort-link">
click here to abort click here to abort
+92
View File
@@ -0,0 +1,92 @@
import { act, renderHook } from '@testing-library/react';
import { describe, expect, it, vi } from 'vitest';
vi.mock('../agent/core/AgentCore', () => ({
AgentCore: {
instance: vi.fn()
}
}));
vi.mock('../utils/chatMessageUtils', () => ({
createStreamingMessage: () => ({
id: 'streaming-message',
role: 'assistant',
content: '<span class="processing-wave">Thinking...</span> click here to abort.',
isStreaming: true,
tokenCount: 0
}),
createMessage: (role: 'user' | 'assistant', content: string) => ({
id: `${role}-message`,
role,
content
})
}));
import { AgentCore } from '../agent/core/AgentCore';
import { useStreamProcessor } from './useStreamProcessor';
import type { ChatMessage } from '../types/projectTypes';
const flushMicrotasks = async (): Promise<void> => {
await Promise.resolve();
await Promise.resolve();
};
describe('useStreamProcessor', () => {
it('switches from Thinking to Processing after the first text token arrives', async () => {
let releaseDone!: () => void;
const doneGate = new Promise<void>((resolve) => {
releaseDone = resolve;
});
vi.spyOn(AgentCore, 'instance').mockReturnValue({
processUserInput: async function* () {
yield { type: 'text', content: 'Hello' };
await doneGate;
yield { type: 'done', content: '' };
}
} as unknown as AgentCore);
const messages = new Map<string, ChatMessage>();
const processingChanges: boolean[] = [];
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: (isProcessing) => {
processingChanges.push(isProcessing);
}
}));
let responsePromise!: Promise<string>;
await act(async () => {
responsePromise = result.current.processStream('test prompt');
await flushMicrotasks();
});
const streamingMessage = [...messages.values()][0];
expect(streamingMessage).toBeDefined();
expect(streamingMessage.content).toContain('<span class="processing-wave">Processing...</span>');
expect(streamingMessage.content).toContain('1 tokens received.');
expect(streamingMessage.content).toContain('click here to abort.');
expect(streamingMessage.tokenCount).toBe(1);
expect(processingChanges).toContain(true);
await act(async () => {
releaseDone();
await responsePromise;
});
expect(processingChanges.at(-1)).toBe(false);
});
});
+3 -1
View File
@@ -16,6 +16,8 @@ interface StreamProcessorResult {
isProcessing: boolean; isProcessing: boolean;
} }
const PROCESSING_WAVE = '<span class="processing-wave">Processing...</span>';
export const useStreamProcessor = (options: StreamProcessorOptions): StreamProcessorResult => { export const useStreamProcessor = (options: StreamProcessorOptions): StreamProcessorResult => {
const { onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange } = options; const { onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange } = options;
const [abortController, setAbortController] = useState<AbortController | null>(null); const [abortController, setAbortController] = useState<AbortController | null>(null);
@@ -56,7 +58,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
onMessageUpdate(currentStreamingId, (msg) => ({ onMessageUpdate(currentStreamingId, (msg) => ({
...msg, ...msg,
content: `<span class="processing-wave">Thinking...</span>${tokenCount > 0 ? ` ${tokenCount} tokens received.` : ''} click here to abort.`, content: `${PROCESSING_WAVE}${tokenCount > 0 ? ` ${tokenCount} tokens received.` : ''} click here to abort.`,
tokenCount tokenCount
})); }));
} else if (chunk.type === 'tool_call' && chunk.toolCall) { } else if (chunk.type === 'tool_call' && chunk.toolCall) {
+11 -3
View File
@@ -17,11 +17,19 @@ vi.mock('../core/KGCore', () => ({
getName: vi.fn().mockReturnValue('Test Project'), getName: vi.fn().mockReturnValue('Test Project'),
getBpm: vi.fn().mockReturnValue(120), getBpm: vi.fn().mockReturnValue(120),
getTimeSignature: vi.fn().mockReturnValue({ numerator: 4, denominator: 4 }), getTimeSignature: vi.fn().mockReturnValue({ numerator: 4, denominator: 4 }),
getTracks: vi.fn().mockReturnValue([]) getTracks: vi.fn().mockReturnValue([]),
getMaxBars: vi.fn().mockReturnValue(64),
getBarWidthMultiplier: vi.fn().mockReturnValue(1),
getIsLooping: vi.fn().mockReturnValue(false),
getLoopingRange: vi.fn().mockReturnValue([0, 0])
}), }),
getSelectedItems: vi.fn().mockReturnValue([]), getSelectedItems: vi.fn().mockReturnValue([]),
setSelectedItems: vi.fn(), setSelectedItems: vi.fn(),
executeCommand: vi.fn() executeCommand: vi.fn(),
setPlayheadUpdateCallback: vi.fn(),
setPlaybackStateChangeCallback: vi.fn(),
setSelectionChangeCallback: vi.fn(),
onSelectionChanged: vi.fn()
}) })
} }
})); }));
@@ -66,4 +74,4 @@ beforeAll(() => {
// Mock URL.createObjectURL (might be needed for file operations) // Mock URL.createObjectURL (might be needed for file operations)
global.URL.createObjectURL = vi.fn(() => 'mocked-url'); global.URL.createObjectURL = vi.fn(() => 'mocked-url');
global.URL.revokeObjectURL = vi.fn(); global.URL.revokeObjectURL = vi.fn();
}); });