feat: Modify loading indicator to support a paused state (#506)

3 files changed, 232 insertions, 36 deletions

diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 8468e61b..45415f39 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -414,15 +414,13 @@ export const useGeminiStream = (
     return StreamProcessingStatus.Completed;
   };
 
-  const streamingState: StreamingState = (() => {
-    if (toolCalls.some((t) => t.status === 'awaiting_approval')) {
-      return StreamingState.WaitingForConfirmation;
-    }
-    if (isResponding || toolCalls.some((t) => t.status === 'executing')) {
-      return StreamingState.Responding;
-    }
-    return StreamingState.Idle;
-  })();
+  const streamingState: StreamingState =
+    isResponding ||
+    toolCalls.some(
+      (t) => t.status === 'awaiting_approval' || t.status === 'executing',
+    )
+      ? StreamingState.Responding
+      : StreamingState.Idle;
 
   const submitQuery = useCallback(
     async (query: PartListUnion) => {
diff --git a/packages/cli/src/ui/hooks/useLoadingIndicator.test.ts b/packages/cli/src/ui/hooks/useLoadingIndicator.test.ts
new file mode 100644
index 00000000..496e13d3
--- /dev/null
+++ b/packages/cli/src/ui/hooks/useLoadingIndicator.test.ts
@@ -0,0 +1,164 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { renderHook, act } from '@testing-library/react';
+import { useLoadingIndicator } from './useLoadingIndicator.js';
+import { StreamingState } from '../types.js';
+import {
+  WITTY_LOADING_PHRASES,
+  PHRASE_CHANGE_INTERVAL_MS,
+} from '../constants.js';
+
+describe('useLoadingIndicator', () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it('should initialize with default values when not responding', () => {
+    const { result } = renderHook(() =>
+      useLoadingIndicator(StreamingState.Idle, false),
+    );
+    expect(result.current.elapsedTime).toBe(0);
+    expect(result.current.currentLoadingPhrase).toBe(WITTY_LOADING_PHRASES[0]);
+    expect(result.current.shouldShowSpinner).toBe(true);
+  });
+
+  describe('when streamingState is Responding', () => {
+    it('should increment elapsedTime and cycle phrases when not paused', () => {
+      const { result } = renderHook(() =>
+        useLoadingIndicator(StreamingState.Responding, false),
+      );
+      expect(result.current.shouldShowSpinner).toBe(true);
+      expect(result.current.currentLoadingPhrase).toBe(
+        WITTY_LOADING_PHRASES[0],
+      );
+
+      act(() => {
+        vi.advanceTimersByTime(1000);
+      });
+      expect(result.current.elapsedTime).toBe(1);
+
+      act(() => {
+        vi.advanceTimersByTime(PHRASE_CHANGE_INTERVAL_MS);
+      });
+      expect(result.current.currentLoadingPhrase).toBe(
+        WITTY_LOADING_PHRASES[1],
+      );
+      expect(result.current.elapsedTime).toBe(
+        1 + PHRASE_CHANGE_INTERVAL_MS / 1000,
+      );
+    });
+
+    it('should pause elapsedTime, show specific phrase, and hide spinner when paused', () => {
+      const { result, rerender } = renderHook(
+        ({ isPaused }) =>
+          useLoadingIndicator(StreamingState.Responding, isPaused),
+        { initialProps: { isPaused: false } },
+      );
+
+      act(() => {
+        vi.advanceTimersByTime(2000);
+      });
+      expect(result.current.elapsedTime).toBe(2);
+      expect(result.current.shouldShowSpinner).toBe(true);
+
+      rerender({ isPaused: true });
+
+      expect(result.current.currentLoadingPhrase).toBe(
+        'Waiting for user confirmation...',
+      );
+      expect(result.current.shouldShowSpinner).toBe(false);
+
+      // Time should not advance while paused
+      const timeBeforePauseAdv = result.current.elapsedTime;
+      act(() => {
+        vi.advanceTimersByTime(3000);
+      });
+      expect(result.current.elapsedTime).toBe(timeBeforePauseAdv);
+
+      // Unpause
+      rerender({ isPaused: false });
+      expect(result.current.shouldShowSpinner).toBe(true);
+      // Phrase should reset to the beginning of witty phrases
+      expect(result.current.currentLoadingPhrase).toBe(
+        WITTY_LOADING_PHRASES[0],
+      );
+
+      act(() => {
+        vi.advanceTimersByTime(1000);
+      });
+      // Elapsed time should resume from where it left off
+      expect(result.current.elapsedTime).toBe(timeBeforePauseAdv + 1);
+    });
+
+    it('should reset timer and phrase when streamingState changes from Responding to Idle', () => {
+      const { result, rerender } = renderHook(
+        ({ streamingState }) => useLoadingIndicator(streamingState, false),
+        { initialProps: { streamingState: StreamingState.Responding } },
+      );
+
+      act(() => {
+        vi.advanceTimersByTime(PHRASE_CHANGE_INTERVAL_MS + 1000);
+      });
+      expect(result.current.elapsedTime).toBe(
+        PHRASE_CHANGE_INTERVAL_MS / 1000 + 1,
+      );
+      expect(result.current.currentLoadingPhrase).toBe(
+        WITTY_LOADING_PHRASES[1],
+      );
+
+      rerender({ streamingState: StreamingState.Idle });
+
+      expect(result.current.elapsedTime).toBe(0);
+      // When idle, the phrase interval should be cleared, but the last phrase might persist
+      // until the next "Responding" state. The important part is that the timer is reset.
+      // Depending on exact timing, it might be the last witty phrase or the first.
+      // For this test, we'll ensure it's one of them.
+      expect(WITTY_LOADING_PHRASES).toContain(
+        result.current.currentLoadingPhrase,
+      );
+    });
+  });
+
+  it('should clear intervals on unmount', () => {
+    const { unmount } = renderHook(() =>
+      useLoadingIndicator(StreamingState.Responding, false),
+    );
+    const clearIntervalSpy = vi.spyOn(global, 'clearInterval');
+    unmount();
+    // Expecting two intervals (elapsedTime and phraseInterval) to be cleared.
+    expect(clearIntervalSpy).toHaveBeenCalledTimes(2);
+  });
+
+  it('should reset to initial witty phrase when unpaused', () => {
+    const { result, rerender } = renderHook(
+      ({ isPaused }) =>
+        useLoadingIndicator(StreamingState.Responding, isPaused),
+      { initialProps: { isPaused: false } },
+    );
+
+    // Advance to the second witty phrase
+    act(() => {
+      vi.advanceTimersByTime(PHRASE_CHANGE_INTERVAL_MS);
+    });
+    expect(result.current.currentLoadingPhrase).toBe(WITTY_LOADING_PHRASES[1]);
+
+    // Pause
+    rerender({ isPaused: true });
+    expect(result.current.currentLoadingPhrase).toBe(
+      'Waiting for user confirmation...',
+    );
+
+    // Unpause
+    rerender({ isPaused: false });
+    expect(result.current.currentLoadingPhrase).toBe(WITTY_LOADING_PHRASES[0]);
+  });
+});
diff --git a/packages/cli/src/ui/hooks/useLoadingIndicator.ts b/packages/cli/src/ui/hooks/useLoadingIndicator.ts
index 6d1d77d4..ac75986a 100644
--- a/packages/cli/src/ui/hooks/useLoadingIndicator.ts
+++ b/packages/cli/src/ui/hooks/useLoadingIndicator.ts
@@ -11,7 +11,10 @@ import {
 } from '../constants.js';
 import { StreamingState } from '../types.js';
 
-export const useLoadingIndicator = (streamingState: StreamingState) => {
+export const useLoadingIndicator = (
+  streamingState: StreamingState,
+  isPaused: boolean,
+) => {
   const [elapsedTime, setElapsedTime] = useState(0);
   const [currentLoadingPhrase, setCurrentLoadingPhrase] = useState(
     WITTY_LOADING_PHRASES[0],
@@ -20,48 +23,79 @@ export const useLoadingIndicator = (streamingState: StreamingState) => {
   const phraseIntervalRef = useRef<NodeJS.Timeout | null>(null);
   const currentPhraseIndexRef = useRef<number>(0);
 
-  // Timer effect for elapsed time during loading
+  const [shouldShowSpinner, setShouldShowSpinner] = useState(true);
+
   useEffect(() => {
     if (streamingState === StreamingState.Responding) {
-      setElapsedTime(0); // Reset timer on new loading start
-      timerRef.current = setInterval(() => {
-        setElapsedTime((prevTime) => prevTime + 1);
-      }, 1000);
-    } else if (timerRef.current) {
-      clearInterval(timerRef.current);
-      timerRef.current = null;
+      if (!isPaused) {
+        if (!timerRef.current) {
+          // No specific action needed here if timer wasn't running and we are not paused.
+          // Elapsed time continues from where it left off or starts from 0 if it's a fresh start.
+        }
+        if (timerRef.current) clearInterval(timerRef.current);
+        timerRef.current = setInterval(() => {
+          setElapsedTime((prevTime) => prevTime + 1);
+        }, 1000);
+      } else {
+        if (timerRef.current) {
+          clearInterval(timerRef.current);
+          timerRef.current = null;
+        }
+      }
+    } else {
+      if (timerRef.current) {
+        clearInterval(timerRef.current);
+        timerRef.current = null;
+      }
+      setElapsedTime(0);
     }
-    // Cleanup on unmount or when isLoading changes
+
     return () => {
       if (timerRef.current) {
         clearInterval(timerRef.current);
+        timerRef.current = null;
       }
     };
-  }, [streamingState]);
+  }, [streamingState, isPaused]);
 
-  // Effect for cycling through witty loading phrases
   useEffect(() => {
     if (streamingState === StreamingState.Responding) {
-      currentPhraseIndexRef.current = 0;
-      setCurrentLoadingPhrase(WITTY_LOADING_PHRASES[0]);
-      phraseIntervalRef.current = setInterval(() => {
-        currentPhraseIndexRef.current =
-          (currentPhraseIndexRef.current + 1) % WITTY_LOADING_PHRASES.length;
-        setCurrentLoadingPhrase(
-          WITTY_LOADING_PHRASES[currentPhraseIndexRef.current],
-        );
-      }, PHRASE_CHANGE_INTERVAL_MS);
-    } else if (phraseIntervalRef.current) {
-      clearInterval(phraseIntervalRef.current);
-      phraseIntervalRef.current = null;
+      if (!isPaused) {
+        setShouldShowSpinner(true);
+        if (!phraseIntervalRef.current) {
+          currentPhraseIndexRef.current = 0;
+          setCurrentLoadingPhrase(WITTY_LOADING_PHRASES[0]);
+          phraseIntervalRef.current = setInterval(() => {
+            currentPhraseIndexRef.current =
+              (currentPhraseIndexRef.current + 1) %
+              WITTY_LOADING_PHRASES.length;
+            setCurrentLoadingPhrase(
+              WITTY_LOADING_PHRASES[currentPhraseIndexRef.current],
+            );
+          }, PHRASE_CHANGE_INTERVAL_MS);
+        }
+      } else {
+        setShouldShowSpinner(false);
+        setCurrentLoadingPhrase('Waiting for user confirmation...');
+        if (phraseIntervalRef.current) {
+          clearInterval(phraseIntervalRef.current);
+          phraseIntervalRef.current = null;
+        }
+      }
+    } else {
+      if (phraseIntervalRef.current) {
+        clearInterval(phraseIntervalRef.current);
+        phraseIntervalRef.current = null;
+      }
     }
-    // Cleanup on unmount or when isLoading changes
+
     return () => {
       if (phraseIntervalRef.current) {
         clearInterval(phraseIntervalRef.current);
+        phraseIntervalRef.current = null;
       }
     };
-  }, [streamingState]);
+  }, [streamingState, isPaused]);
 
-  return { elapsedTime, currentLoadingPhrase };
+  return { elapsedTime, currentLoadingPhrase, shouldShowSpinner };
 };