From 25e48185ea92f510783ea2011773c79b2c486d95 Mon Sep 17 00:00:00 2001
From: Victoria Wang <quan.wang147@gmail.com>
Date: Wed, 13 May 2026 23:16:50 -0400
Subject: [PATCH] fix: prevent Recent Token Usage from being cleared during
 gateway status reconciliation

---
 src/pages/Models/index.tsx    | 45 ++++++++++++++++++++---------------
 tests/e2e/token-usage.spec.ts | 23 ++++++++++++++++++
 2 files changed, 49 insertions(+), 19 deletions(-)

diff --git a/src/pages/Models/index.tsx b/src/pages/Models/index.tsx
index 5ee373c..7b6dcae 100644
--- a/src/pages/Models/index.tsx
+++ b/src/pages/Models/index.tsx
@@ -28,7 +28,6 @@ export function Models() {
   const { t } = useTranslation(['dashboard', 'settings']);
   const gatewayStatus = useGatewayStore((state) => state.status);
   const devModeUnlocked = useSettingsStore((state) => state.devModeUnlocked);
-  const isGatewayRunning = gatewayStatus.state === 'running';
   const usageFetchMaxAttempts = getRendererPlatform() === 'win32'
     ? WINDOWS_USAGE_FETCH_MAX_ATTEMPTS
     : DEFAULT_USAGE_FETCH_MAX_ATTEMPTS;
@@ -76,7 +75,7 @@ export function Models() {
   type FetchAction =
     | { type: 'start' }
     | { type: 'done'; data: UsageHistoryEntry[] }
-    | { type: 'reset' };
+    | { type: 'idle' };
 
   const [fetchState, dispatchFetch] = useReducer(
     (state: FetchState, action: FetchAction): FetchState => {
@@ -85,8 +84,8 @@ export function Models() {
           return { status: 'loading', data: state.data };
         case 'done':
           return { status: 'done', data: action.data };
-        case 'reset':
-          return { status: 'idle', data: [] };
+        case 'idle':
+          return { status: 'idle', data: state.data };
         default:
           return state;
       }
@@ -96,6 +95,11 @@ export function Models() {
 
   const usageFetchTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const usageFetchGenerationRef = useRef(0);
+  const usageDataRef = useRef<UsageHistoryEntry[]>([]);
+
+  useEffect(() => {
+    usageDataRef.current = fetchState.data;
+  }, [fetchState.data]);
 
   useEffect(() => {
     trackUiEvent('models.page_viewed');
@@ -107,11 +111,6 @@ export function Models() {
       usageFetchTimerRef.current = null;
     }
 
-    if (!isGatewayRunning) {
-      dispatchFetch({ type: 'reset' });
-      return;
-    }
-
     dispatchFetch({ type: 'start' });
     const generation = usageFetchGenerationRef.current + 1;
     usageFetchGenerationRef.current = generation;
@@ -129,7 +128,10 @@ export function Models() {
         generation,
         restartMarker,
       });
-      dispatchFetch({ type: 'done', data: [] });
+      dispatchFetch({
+        type: 'done',
+        data: usageDataRef.current,
+      });
     }, 30_000);
 
     const fetchUsageHistoryWithRetry = async (attempt: number) => {
@@ -170,7 +172,10 @@ export function Models() {
               restartMarker,
             });
           }
-          dispatchFetch({ type: 'done', data: normalized });
+          dispatchFetch({
+            type: 'done',
+            data: normalized.length > 0 ? normalized : usageDataRef.current,
+          });
         }
       } catch (error) {
         if (usageFetchGenerationRef.current !== generation) return;
@@ -192,7 +197,10 @@ export function Models() {
           }, USAGE_FETCH_RETRY_DELAY_MS);
           return;
         }
-        dispatchFetch({ type: 'done', data: [] });
+        dispatchFetch({
+          type: 'done',
+          data: usageDataRef.current,
+        });
         trackUiEvent('models.token_usage_fetch_exhausted', {
           generation,
           attempt,
@@ -211,13 +219,12 @@ export function Models() {
         usageFetchTimerRef.current = null;
       }
     };
-  }, [isGatewayRunning, gatewayStatus.connectedAt, gatewayStatus.pid, usageFetchMaxAttempts]);
+  }, [gatewayStatus.connectedAt, gatewayStatus.pid, usageFetchMaxAttempts]);
 
-  const visibleUsageHistory = useMemo(() => (
-    isGatewayRunning
-      ? fetchState.data.filter((entry) => !shouldHideUsageEntry(entry))
-      : []
-  ), [fetchState.data, isGatewayRunning]);
+  const visibleUsageHistory = useMemo(
+    () => fetchState.data.filter((entry) => !shouldHideUsageEntry(entry)),
+    [fetchState.data],
+  );
 
   const filteredUsageHistory = useMemo(
     () => filterUsageHistoryByWindow(visibleUsageHistory, usageWindow),
@@ -236,7 +243,7 @@ export function Models() {
     () => filteredUsageHistory.slice((safeUsagePage - 1) * usagePageSize, safeUsagePage * usagePageSize),
     [filteredUsageHistory, safeUsagePage],
   );
-  const usageLoading = isGatewayRunning && fetchState.status === 'loading';
+  const usageLoading = fetchState.status === 'loading';
 
   return (
     <div data-testid="models-page" className="page-view">
diff --git a/tests/e2e/token-usage.spec.ts b/tests/e2e/token-usage.spec.ts
index 3b7f313..257b6cf 100644
--- a/tests/e2e/token-usage.spec.ts
+++ b/tests/e2e/token-usage.spec.ts
@@ -195,4 +195,27 @@ test.describe('InvestClaw token usage history', () => {
     await expect(page.locator('[data-testid="token-usage-entry"]', { hasText: GATEWAY_INJECTED_SESSION_ID })).toHaveCount(0);
     await expect(page.locator('[data-testid="token-usage-entry"]', { hasText: DELIVERY_MIRROR_SESSION_ID })).toHaveCount(0);
   });
+
+  test('keeps recent token usage visible if gateway stops after the models page has loaded', async ({ page, homeDir }) => {
+    await seedTokenUsageTranscripts(homeDir);
+    await completeSetup(page);
+    await waitForGatewayRunning(page);
+    await validateUsageHistory(page);
+    await page.getByTestId('sidebar-nav-models').click();
+    await expect(page.getByTestId('models-page')).toBeVisible();
+
+    const usageEntryRows = page.getByTestId('token-usage-entry');
+    await expect.poll(async () => await usageEntryRows.count()).toBe(2);
+
+    await page.evaluate(async () => {
+      await window.electron.ipcRenderer.invoke('gateway:stop');
+    });
+
+    await expect.poll(async () => {
+      const status = await page.evaluate(async () => window.electron.ipcRenderer.invoke('gateway:status'));
+      return status?.state;
+    }).toBe('stopped');
+
+    await expect.poll(async () => await usageEntryRows.count()).toBe(2);
+  });
 });