Skip to content

fix: prevent Recent Token Usage from being cleared during gateway status reconciliation #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SmartBunBun wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: prevent Recent Token Usage from being cleared during gateway status reconciliation #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 26 additions & 19 deletions src/pages/Models/index.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 => {
Expand All @@ -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;
}
Expand All @@ -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');
Expand All @@ -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;
Expand All @@ -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) => {
Expand Down Expand Up @@ -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;
Expand All @@ -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,
Expand All @@ -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),
Expand All @@ -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">
Expand Down
23 changes: 23 additions & 0 deletions tests/e2e/token-usage.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});