From 2bd6925f27477a684ced909c6fa443cd562d56d4 Mon Sep 17 00:00:00 2001
From: Xiaohan-Tian <xiaohan.tian@gmail.com>
Date: Thu, 14 May 2026 21:37:06 -0700
Subject: [PATCH] fix: switch to processing status when receiving response
 tokens from LLM

---
 src/components/chat/AssistantMessage.test.tsx | 25 +++++
 src/components/chat/AssistantMessage.tsx      | 12 ++-
 src/hooks/useStreamProcessor.test.ts          | 92 +++++++++++++++++++
 src/hooks/useStreamProcessor.ts               |  4 +-
 src/test/setup.ts                             | 14 ++-
 5 files changed, 139 insertions(+), 8 deletions(-)
 create mode 100644 src/components/chat/AssistantMessage.test.tsx
 create mode 100644 src/hooks/useStreamProcessor.test.ts

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([
+    '<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);
+  });
+});
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<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
diff --git a/src/hooks/useStreamProcessor.test.ts b/src/hooks/useStreamProcessor.test.ts
new file mode 100644
index 0000000..8685d2d
--- /dev/null
+++ b/src/hooks/useStreamProcessor.test.ts
@@ -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);
+  });
+});
diff --git a/src/hooks/useStreamProcessor.ts b/src/hooks/useStreamProcessor.ts
index 3a6c210..45209d4 100644
--- a/src/hooks/useStreamProcessor.ts
+++ b/src/hooks/useStreamProcessor.ts
@@ -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) {
diff --git a/src/test/setup.ts b/src/test/setup.ts
index d1aba87..ad6a39a 100644
--- a/src/test/setup.ts
+++ b/src/test/setup.ts
@@ -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();
-});
\ No newline at end of file
+});