diff --git a/src/components/chat/AssistantMessage.test.tsx b/src/components/chat/AssistantMessage.test.tsx
new file mode 100644
index 0000000..04d68b2
--- /dev/null
+++ b/src/components/chat/AssistantMessage.test.tsx
@@ -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([
+ 'Thinking... click here to abort.',
+ 'Processing... 3 tokens received. click here to abort.'
+ ])('renders the abort action for streaming status content: %s', (content) => {
+ const onAbort = vi.fn();
+
+ render(
+
+ );
+
+ const abortButton = screen.getByRole('button', { name: 'click here to abort' });
+ expect(abortButton).toBeInTheDocument();
+ fireEvent.click(abortButton);
+ expect(onAbort).toHaveBeenCalledTimes(1);
+ });
+});
diff --git a/src/components/chat/AssistantMessage.tsx b/src/components/chat/AssistantMessage.tsx
index ad196be..96b2b93 100644
--- a/src/components/chat/AssistantMessage.tsx
+++ b/src/components/chat/AssistantMessage.tsx
@@ -44,22 +44,26 @@ const AssistantMessage: React.FC = ({ 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('Thinking...');
+ const processingWaveMarkup = processingWaveLabels
+ .map(label => `${label}`)
+ .find(markup => content.includes(markup));
- if (hasProcessingWave) {
+ if (processingWaveMarkup) {
const parts = content.split('click here to abort');
const beforeAbort = parts[0].replace(
- 'Thinking...',
+ processingWaveMarkup,
''
);
+ const waveLabel = processingWaveLabels.find(label => processingWaveMarkup.includes(label)) ?? 'Thinking...';
return (
- Thinking...
+ {waveLabel}
{beforeAbort}