fix: switch to processing status when receiving response tokens from LLM
This commit is contained in:
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -44,22 +44,26 @@ const AssistantMessage: React.FC<AssistantMessageProps> = ({ content, isStreamin
|
||||
const prefillTps = formatTps(performanceInfo?.prefillTps);
|
||||
const generationTps = formatTps(performanceInfo?.generationTps);
|
||||
const hasPerformanceInfo = Boolean(prefillTps || generationTps);
|
||||
const processingWaveLabels = ['Thinking...', 'Processing...'];
|
||||
|
||||
const renderContent = () => {
|
||||
// Handle special abort link for streaming messages
|
||||
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 beforeAbort = parts[0].replace(
|
||||
'<span class="processing-wave">Thinking...</span>',
|
||||
processingWaveMarkup,
|
||||
''
|
||||
);
|
||||
const waveLabel = processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? 'Thinking...';
|
||||
|
||||
return (
|
||||
<span>
|
||||
<span className="processing-wave">Thinking...</span>
|
||||
<span className="processing-wave">{waveLabel}</span>
|
||||
{beforeAbort}
|
||||
<button onClick={onAbort} className="abort-link">
|
||||
click here to abort
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -16,6 +16,8 @@ interface StreamProcessorResult {
|
||||
isProcessing: boolean;
|
||||
}
|
||||
|
||||
const PROCESSING_WAVE = '<span class="processing-wave">Processing...</span>';
|
||||
|
||||
export const useStreamProcessor = (options: StreamProcessorOptions): StreamProcessorResult => {
|
||||
const { onMessageUpdate, onMessageAdd, onMessageRemove, onProcessingChange } = options;
|
||||
const [abortController, setAbortController] = useState<AbortController | null>(null);
|
||||
@@ -56,7 +58,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce
|
||||
|
||||
onMessageUpdate(currentStreamingId, (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
|
||||
}));
|
||||
} else if (chunk.type === 'tool_call' && chunk.toolCall) {
|
||||
|
||||
+11
-3
@@ -17,11 +17,19 @@ vi.mock('../core/KGCore', () => ({
|
||||
getName: vi.fn().mockReturnValue('Test Project'),
|
||||
getBpm: vi.fn().mockReturnValue(120),
|
||||
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([]),
|
||||
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)
|
||||
global.URL.createObjectURL = vi.fn(() => 'mocked-url');
|
||||
global.URL.revokeObjectURL = vi.fn();
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user