Skip to content

Commit 23da875

Browse files
authored
Merge pull request #237 from seakee/fix/codex-quota-cache
Fix manual Codex quota refresh persistence
2 parents d293c0c + 8fccc30 commit 23da875

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)