Skip to content

Commit 8fccc30

Browse files
committed
🐛 fix(web): persist manual Codex quota refreshes
Persist successful manual Codex quota refresh results and reuse them after full page reloads. Share quota display resolution between quota management and auth files so stale header snapshots no longer override newer manual data. Clear cached quota state on auth identity changes, logout, and login storage cleanup to avoid cross-account leakage.
1 parent d293c0c commit 8fccc30

10 files changed

Lines changed: 278 additions & 50 deletions

File tree

apps/web/src/components/quota/QuotaSection.tsx

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import {
1818
import { QuotaCard } from './QuotaCard';
1919
import type { QuotaStatusState } from './QuotaCard';
2020
import { useQuotaLoader } from './useQuotaLoader';
21-
import type { QuotaConfig, QuotaSortMode } from './quotaConfigs';
21+
import { resolveQuotaDisplayState, type QuotaConfig, type QuotaSortMode } from './quotaConfigs';
2222
import { resolveQuotaAccountDisplayText } from './quotaDisplay';
2323
import {
2424
DEFAULT_QUOTA_ACCOUNT_DISPLAY_MODE,
@@ -192,21 +192,12 @@ export function QuotaSection<TState extends QuotaStatusState, TData>({
192192
const getDisplayQuota = useCallback(
193193
(file: AuthFileItem): TState | undefined => {
194194
const activeQuota = quota[file.name];
195-
if (activeQuota && activeQuota.status !== 'idle' && activeQuota.status !== 'error') {
196-
return activeQuota;
197-
}
198-
if (
199-
activeQuota?.status === 'error' &&
200-
(activeQuota as { errorStatus?: number | null }).errorStatus === 401
201-
) {
202-
return activeQuota;
203-
}
204195
const observedQuota = config.buildObservedState?.(
205196
file,
206197
getHighConfidenceUsageHeaderSnapshotForAuthFile(headerSnapshotLookup, file),
207198
t
208199
);
209-
return observedQuota ?? activeQuota;
200+
return resolveQuotaDisplayState(activeQuota, observedQuota);
210201
},
211202
[config, headerSnapshotLookup, quota, t]
212203
);

apps/web/src/components/quota/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,5 +12,6 @@ export {
1212
KIMI_CONFIG,
1313
XAI_CONFIG,
1414
buildObservedCodexQuotaState,
15+
resolveQuotaDisplayState,
1516
} from './quotaConfigs';
1617
export type { QuotaConfig } from './quotaConfigs';
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import { describe, expect, it } from 'vitest';
2+
import { resolveQuotaDisplayState } from './quotaConfigs';
3+
4+
type TestQuotaState = {
5+
status: 'idle' | 'loading' | 'success' | 'error';
6+
errorStatus?: number;
7+
fetchedAtMs?: number;
8+
observedAtMs?: number;
9+
observedFromUsageHeaders?: boolean;
10+
windows?: unknown[];
11+
};
12+
13+
describe('resolveQuotaDisplayState', () => {
14+
it('keeps a newer manual quota refresh over an older header snapshot', () => {
15+
const activeQuota: TestQuotaState = {
16+
status: 'success',
17+
fetchedAtMs: 2_000,
18+
windows: [],
19+
};
20+
const observedQuota: TestQuotaState = {
21+
status: 'success',
22+
observedAtMs: 1_000,
23+
observedFromUsageHeaders: true,
24+
windows: [],
25+
};
26+
27+
expect(resolveQuotaDisplayState(activeQuota, observedQuota)).toBe(activeQuota);
28+
});
29+
30+
it('uses a newer header snapshot when it is fresher than the manual quota refresh', () => {
31+
const activeQuota: TestQuotaState = {
32+
status: 'success',
33+
fetchedAtMs: 1_000,
34+
windows: [],
35+
};
36+
const observedQuota: TestQuotaState = {
37+
status: 'success',
38+
observedAtMs: 2_000,
39+
observedFromUsageHeaders: true,
40+
windows: [],
41+
};
42+
43+
expect(resolveQuotaDisplayState(activeQuota, observedQuota)).toBe(observedQuota);
44+
});
45+
46+
it('keeps 401 quota errors so reauth controls stay visible', () => {
47+
const activeQuota: TestQuotaState = {
48+
status: 'error',
49+
errorStatus: 401,
50+
};
51+
const observedQuota: TestQuotaState = {
52+
status: 'success',
53+
observedAtMs: 2_000,
54+
};
55+
56+
expect(resolveQuotaDisplayState(activeQuota, observedQuota)).toBe(activeQuota);
57+
});
58+
});

apps/web/src/components/quota/quotaConfigs.ts

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,38 @@ const getCodexSearchText = (
230230
];
231231
};
232232

233+
type DisplayQuotaState = {
234+
status?: 'idle' | 'loading' | 'success' | 'error';
235+
errorStatus?: number | null;
236+
fetchedAtMs?: number;
237+
observedAtMs?: number;
238+
};
239+
240+
const readFiniteTimestamp = (value: unknown): number | null =>
241+
typeof value === 'number' && Number.isFinite(value) ? value : null;
242+
243+
export const resolveQuotaDisplayState = <TState extends DisplayQuotaState>(
244+
activeQuota: TState | undefined,
245+
observedQuota: TState | undefined
246+
): TState | undefined => {
247+
if (activeQuota && activeQuota.status !== 'idle' && activeQuota.status !== 'error') {
248+
if (activeQuota.status === 'success' && observedQuota?.status === 'success') {
249+
const fetchedAtMs = readFiniteTimestamp(activeQuota.fetchedAtMs);
250+
const observedAtMs = readFiniteTimestamp(observedQuota.observedAtMs);
251+
if (fetchedAtMs !== null && observedAtMs !== null && observedAtMs > fetchedAtMs) {
252+
return observedQuota;
253+
}
254+
}
255+
return activeQuota;
256+
}
257+
258+
if (activeQuota?.status === 'error' && activeQuota.errorStatus === 401) {
259+
return activeQuota;
260+
}
261+
262+
return observedQuota ?? activeQuota;
263+
};
264+
233265
export const buildObservedCodexQuotaState = (
234266
file: AuthFileItem,
235267
snapshot: UsageHeaderSnapshot | undefined,

apps/web/src/features/authFiles/AuthFilesPage.tsx

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ import { Select } from '@/components/ui/Select';
2424
import { IconFilterAll, IconSearch } from '@/components/ui/icons';
2525
import { EmptyState } from '@/components/ui/EmptyState';
2626
import { ToggleSwitch } from '@/components/ui/ToggleSwitch';
27-
import { buildObservedCodexQuotaState } from '@/components/quota';
27+
import { buildObservedCodexQuotaState, resolveQuotaDisplayState } from '@/components/quota';
2828
import { copyToClipboard } from '@/utils/clipboard';
2929
import { resolveAuthProvider } from '@/utils/quota';
3030
import {
@@ -645,18 +645,12 @@ export function AuthFilesPage() {
645645
(file: AuthFileItem): CodexQuotaState | undefined => {
646646
if (resolveAuthProvider(file) !== 'codex') return undefined;
647647
const activeQuota = codexQuota[file.name];
648-
if (activeQuota && activeQuota.status !== 'idle' && activeQuota.status !== 'error') {
649-
return activeQuota;
650-
}
651-
if (activeQuota?.status === 'error' && activeQuota.errorStatus === 401) {
652-
return activeQuota;
653-
}
654648
const observedQuota = buildObservedCodexQuotaState(
655649
file,
656650
getHighConfidenceUsageHeaderSnapshotForAuthFile(headerSnapshotLookup, file),
657651
t
658652
);
659-
return observedQuota ?? activeQuota;
653+
return resolveQuotaDisplayState(activeQuota, observedQuota);
660654
},
661655
[codexQuota, headerSnapshotLookup, t]
662656
);

apps/web/src/features/system/SystemPage.tsx

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import {
1212
} from '@/stores';
1313
import { apiKeysApi } from '@/services/api/apiKeys';
1414
import { classifyModels } from '@/utils/models';
15-
import { STORAGE_KEY_AUTH } from '@/utils/constants';
15+
import { STORAGE_KEY_AUTH, STORAGE_KEY_QUOTA_CACHE } from '@/utils/constants';
1616
import iconGemini from '@/assets/icons/gemini.svg';
1717
import iconClaude from '@/assets/icons/claude.svg';
1818
import iconOpenaiLight from '@/assets/icons/openai-light.svg';
@@ -167,7 +167,14 @@ export function SystemPage() {
167167
onConfirm: () => {
168168
auth.logout();
169169
if (typeof localStorage === 'undefined') return;
170-
const keysToRemove = [STORAGE_KEY_AUTH, 'isLoggedIn', 'apiBase', 'apiUrl', 'managementKey'];
170+
const keysToRemove = [
171+
STORAGE_KEY_AUTH,
172+
STORAGE_KEY_QUOTA_CACHE,
173+
'isLoggedIn',
174+
'apiBase',
175+
'apiUrl',
176+
'managementKey',
177+
];
171178
keysToRemove.forEach((key) => localStorage.removeItem(key));
172179
showNotification(t('notification.login_storage_cleared'), 'success');
173180
},

apps/web/src/stores/useAuthStore.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import { apiClient } from '@/services/api/client';
1919
import { usageServiceApi } from '@/services/api/usageService';
2020
import { useConfigStore } from './useConfigStore';
2121
import { useModelsStore } from './useModelsStore';
22+
import { useQuotaStore } from './useQuotaStore';
2223
import { useUsageServiceStore } from './useUsageServiceStore';
2324
import { detectApiBaseFromLocation, normalizeApiBase } from '@/utils/connection';
2425

@@ -162,8 +163,16 @@ export const useAuthStore = create<AuthStoreState>()(
162163
const rememberPassword = credentials.rememberPassword ?? get().rememberPassword ?? false;
163164
const sessionMode = credentials.sessionMode ?? get().sessionMode;
164165
const sessionPanelBase = normalizeApiBase(credentials.sessionPanelBase || get().sessionPanelBase);
166+
const previousApiBase = get().apiBase;
167+
const previousManagementKey = get().managementKey;
168+
const shouldClearQuotaCache =
169+
Boolean(previousApiBase || previousManagementKey) &&
170+
(previousApiBase !== apiBase || previousManagementKey !== managementKey);
165171

166172
const markAuthenticated = (result: LoginResult = {}) => {
173+
if (shouldClearQuotaCache) {
174+
useQuotaStore.getState().clearQuotaCache();
175+
}
167176
apiClient.setConfig({ apiBase, managementKey });
168177
set({
169178
isAuthenticated: true,
@@ -238,6 +247,7 @@ export const useAuthStore = create<AuthStoreState>()(
238247
restoreSessionPromise = null;
239248
useConfigStore.getState().clearCache();
240249
useModelsStore.getState().clearCache();
250+
useQuotaStore.getState().clearQuotaCache();
241251
useUsageServiceStore.getState().clearUsageServiceConfig();
242252
apiClient.setConfig({ apiBase: '', managementKey: '' });
243253
set({
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
2+
import type { CodexQuotaState } from '@/types';
3+
4+
type StorageLike = {
5+
getItem: (key: string) => string | null;
6+
setItem: (key: string, value: string) => void;
7+
removeItem: (key: string) => void;
8+
clear: () => void;
9+
};
10+
11+
const createMemoryStorage = (): StorageLike => {
12+
const store = new Map<string, string>();
13+
return {
14+
getItem: (key) => (store.has(key) ? (store.get(key) as string) : null),
15+
setItem: (key, value) => {
16+
store.set(key, value);
17+
},
18+
removeItem: (key) => {
19+
store.delete(key);
20+
},
21+
clear: () => {
22+
store.clear();
23+
},
24+
};
25+
};
26+
27+
const readPersistedCodexQuota = async () => {
28+
const { STORAGE_KEY_QUOTA_CACHE } = await import('@/utils/constants');
29+
const { obfuscatedStorage } = await import('@/services/storage/secureStorage');
30+
const persisted = obfuscatedStorage.getItem<{
31+
state?: { codexQuota?: Record<string, CodexQuotaState> };
32+
}>(STORAGE_KEY_QUOTA_CACHE);
33+
return persisted?.state?.codexQuota ?? {};
34+
};
35+
36+
describe('useQuotaStore persistence', () => {
37+
let storage: StorageLike;
38+
39+
beforeEach(() => {
40+
vi.resetModules();
41+
storage = createMemoryStorage();
42+
vi.stubGlobal('localStorage', storage);
43+
});
44+
45+
afterEach(() => {
46+
vi.unstubAllGlobals();
47+
});
48+
49+
it('persists only manually fetched Codex success states', async () => {
50+
const { useQuotaStore } = await import('./useQuotaStore');
51+
52+
useQuotaStore.getState().setCodexQuota({
53+
manual: {
54+
status: 'success',
55+
windows: [],
56+
fetchedAtMs: 2_000,
57+
},
58+
observed: {
59+
status: 'success',
60+
windows: [],
61+
observedFromUsageHeaders: true,
62+
observedAtMs: 1_000,
63+
},
64+
failed: {
65+
status: 'error',
66+
windows: [],
67+
error: 'failed',
68+
},
69+
});
70+
71+
expect(Object.keys(await readPersistedCodexQuota())).toEqual(['manual']);
72+
});
73+
74+
it('clears quota state and persisted quota cache together', async () => {
75+
const { useQuotaStore } = await import('./useQuotaStore');
76+
77+
useQuotaStore.getState().setCodexQuota({
78+
manual: {
79+
status: 'success',
80+
windows: [],
81+
fetchedAtMs: 2_000,
82+
},
83+
});
84+
85+
useQuotaStore.getState().clearQuotaCache();
86+
87+
expect(useQuotaStore.getState().codexQuota).toEqual({});
88+
expect(await readPersistedCodexQuota()).toEqual({});
89+
});
90+
});

0 commit comments

Comments
 (0)