Skip to content

Commit 3afc5d7

Browse files
ndycodendycode
authored
fix(codex-limits): harden usage fetch and dedup refresh
Harden codex-limits usage fetching, refresh propagation, and deduplication semantics; close the PR 70 review loop with transactional persistence, token redaction, and targeted regression coverage.
1 parent 155eba0 commit 3afc5d7

File tree

2 files changed

+819
-79
lines changed

2 files changed

+819
-79
lines changed

‎index.ts‎

Lines changed: 204 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -4298,6 +4298,125 @@ while (attempted.size < Math.max(1, accountCount)) {
42984298
return name.replace(/[_-]+/g, " ").replace(/\b\w/g, (match) => match.toUpperCase());
42994299
};
43004300

4301+
const sanitizeUsageErrorMessage = (status: number, bodyText: string): string => {
4302+
const normalized = bodyText.replace(/\s+/g, " ").trim();
4303+
const redacted = normalized
4304+
.replace(/Bearer\s+\S+/gi, "Bearer [redacted]")
4305+
.replace(/eyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+/g, "[redacted-token]")
4306+
.replace(/\bsk-[A-Za-z0-9][A-Za-z0-9._:-]{19,}\b/gi, "[redacted-token]")
4307+
.replace(/\b[a-f0-9]{40,}\b/gi, "[redacted-token]");
4308+
return redacted ? `HTTP ${status}: ${redacted.slice(0, 200)}` : `HTTP ${status}`;
4309+
};
4310+
4311+
const isAbortError = (error: unknown): boolean =>
4312+
(error instanceof Error && error.name === "AbortError") ||
4313+
(typeof DOMException !== "undefined" && error instanceof DOMException && error.name === "AbortError");
4314+
4315+
const applyRefreshedCredentials = (
4316+
target: {
4317+
refreshToken: string;
4318+
accessToken?: string;
4319+
expiresAt?: number;
4320+
},
4321+
result: {
4322+
refresh: string;
4323+
access: string;
4324+
expires: number;
4325+
},
4326+
): void => {
4327+
target.refreshToken = result.refresh;
4328+
target.accessToken = result.access;
4329+
target.expiresAt = result.expires;
4330+
};
4331+
const usageErrorBodyMaxChars = 4096;
4332+
4333+
const persistRefreshedCredentials = async (params: {
4334+
previousRefreshToken: string;
4335+
accountId?: string;
4336+
organizationId?: string;
4337+
email?: string;
4338+
refreshResult: {
4339+
refresh: string;
4340+
access: string;
4341+
expires: number;
4342+
};
4343+
}): Promise<boolean> => {
4344+
return await withAccountStorageTransaction(async (current, persist) => {
4345+
const latestStorage: AccountStorageV3 =
4346+
current ??
4347+
({
4348+
version: 3,
4349+
accounts: [],
4350+
activeIndex: 0,
4351+
activeIndexByFamily: {},
4352+
} satisfies AccountStorageV3);
4353+
4354+
const uniqueMatch = <T>(matches: T[]): T | undefined =>
4355+
matches.length === 1 ? matches[0] : undefined;
4356+
4357+
let updated = false;
4358+
if (params.previousRefreshToken) {
4359+
for (const storedAccount of latestStorage.accounts) {
4360+
if (storedAccount.refreshToken === params.previousRefreshToken) {
4361+
applyRefreshedCredentials(storedAccount, params.refreshResult);
4362+
updated = true;
4363+
}
4364+
}
4365+
}
4366+
4367+
if (!updated) {
4368+
const normalizedOrganizationId = params.organizationId?.trim() ?? "";
4369+
const normalizedEmail = params.email?.trim().toLowerCase();
4370+
const orgScopedMatches = params.accountId
4371+
? latestStorage.accounts.filter(
4372+
(storedAccount) =>
4373+
storedAccount.accountId === params.accountId &&
4374+
(storedAccount.organizationId?.trim() ?? "") === normalizedOrganizationId,
4375+
)
4376+
: [];
4377+
const accountIdMatches = params.accountId
4378+
? latestStorage.accounts.filter(
4379+
(storedAccount) => storedAccount.accountId === params.accountId,
4380+
)
4381+
: [];
4382+
const emailMatches =
4383+
normalizedEmail && !params.accountId
4384+
? latestStorage.accounts.filter(
4385+
(storedAccount) =>
4386+
storedAccount.email?.trim().toLowerCase() === normalizedEmail,
4387+
)
4388+
: [];
4389+
4390+
const fallbackTarget =
4391+
uniqueMatch(orgScopedMatches) ??
4392+
uniqueMatch(accountIdMatches) ??
4393+
uniqueMatch(emailMatches);
4394+
4395+
if (fallbackTarget) {
4396+
applyRefreshedCredentials(fallbackTarget, params.refreshResult);
4397+
updated = true;
4398+
}
4399+
}
4400+
4401+
if (updated) {
4402+
await persist(latestStorage);
4403+
}
4404+
if (!updated) {
4405+
logWarn(
4406+
`[${PLUGIN_NAME}] persistRefreshedCredentials could not find a matching stored account. Refreshed credentials remain in-memory for this invocation only.`,
4407+
{
4408+
accountId: params.accountId,
4409+
organizationId: params.organizationId,
4410+
},
4411+
);
4412+
}
4413+
4414+
return updated;
4415+
});
4416+
};
4417+
4418+
const usageFetchTimeoutMs = getFetchTimeoutMs(loadPluginConfig());
4419+
43014420
const fetchUsage = async (params: {
43024421
accountId: string;
43034422
accessToken: string;
@@ -4307,16 +4426,39 @@ while (attempted.size < Math.max(1, accountCount)) {
43074426
organizationId: params.organizationId,
43084427
});
43094428
headers.set("accept", "application/json");
4429+
const controller = new AbortController();
4430+
const timeout = setTimeout(() => controller.abort(), usageFetchTimeoutMs);
43104431

4311-
const response = await fetch(`${CODEX_BASE_URL}/wham/usage`, {
4312-
method: "GET",
4313-
headers,
4314-
});
4315-
if (!response.ok) {
4316-
const bodyText = await response.text().catch(() => "");
4317-
throw new Error(bodyText || `HTTP ${response.status}`);
4432+
try {
4433+
const response = await fetch(`${CODEX_BASE_URL}/wham/usage`, {
4434+
method: "GET",
4435+
headers,
4436+
signal: controller.signal,
4437+
});
4438+
if (!response.ok) {
4439+
let bodyText = "";
4440+
try {
4441+
bodyText = (await response.text()).slice(0, usageErrorBodyMaxChars);
4442+
} catch (error) {
4443+
if (isAbortError(error) || controller.signal.aborted) {
4444+
throw new Error("Usage request timed out");
4445+
}
4446+
throw error;
4447+
}
4448+
if (controller.signal.aborted) {
4449+
throw new Error("Usage request timed out");
4450+
}
4451+
throw new Error(sanitizeUsageErrorMessage(response.status, bodyText));
4452+
}
4453+
return (await response.json()) as UsagePayload;
4454+
} catch (error) {
4455+
if (isAbortError(error)) {
4456+
throw new Error("Usage request timed out");
4457+
}
4458+
throw error;
4459+
} finally {
4460+
clearTimeout(timeout);
43184461
}
4319-
return (await response.json()) as UsagePayload;
43204462
};
43214463

43224464
// Deduplicate accounts by refreshToken (same credential = same limits)
@@ -4325,22 +4467,40 @@ while (attempted.size < Math.max(1, accountCount)) {
43254467
for (let i = 0; i < storage.accounts.length; i++) {
43264468
const acct = storage.accounts[i];
43274469
if (!acct) continue;
4328-
if (seenTokens.has(acct.refreshToken)) continue;
4329-
seenTokens.add(acct.refreshToken);
4470+
const refreshToken =
4471+
typeof acct.refreshToken === "string" ? acct.refreshToken.trim() : "";
4472+
if (refreshToken && seenTokens.has(refreshToken)) continue;
4473+
if (refreshToken) seenTokens.add(refreshToken);
43304474
uniqueIndices.push(i);
43314475
}
43324476

43334477
const lines: string[] = ui.v2Enabled
43344478
? [...formatUiHeader(ui, "Codex limits"), ""]
43354479
: [`Codex limits (${uniqueIndices.length} account${uniqueIndices.length === 1 ? "" : "s"}):`, ""];
43364480
const activeIndex = resolveActiveIndex(storage, "codex");
4481+
const activeRefreshToken =
4482+
typeof activeIndex === "number" && activeIndex >= 0 && activeIndex < storage.accounts.length
4483+
? storage.accounts[activeIndex]?.refreshToken?.trim() || undefined
4484+
: undefined;
43374485
let storageChanged = false;
43384486

43394487
for (const i of uniqueIndices) {
43404488
const account = storage.accounts[i];
43414489
if (!account) continue;
4342-
const label = formatCommandAccountLabel(account, i);
4343-
const activeSuffix = i === activeIndex ? (ui.v2Enabled ? ` ${formatUiBadge(ui, "active", "accent")}` : " [active]") : "";
4490+
const sharesActiveCredential =
4491+
!!activeRefreshToken && account.refreshToken === activeRefreshToken;
4492+
const displayIndex =
4493+
sharesActiveCredential && typeof activeIndex === "number" ? activeIndex : i;
4494+
const displayAccount = storage.accounts[displayIndex];
4495+
if (sharesActiveCredential && !displayAccount) {
4496+
logWarn(
4497+
`[${PLUGIN_NAME}] active account entry missing for index ${displayIndex}, falling back to account ${i}`,
4498+
);
4499+
}
4500+
const effectiveDisplayAccount = displayAccount ?? account;
4501+
const label = formatCommandAccountLabel(effectiveDisplayAccount, displayIndex);
4502+
const isActive = i === activeIndex || sharesActiveCredential;
4503+
const activeSuffix = isActive ? (ui.v2Enabled ? ` ${formatUiBadge(ui, "active", "accent")}` : " [active]") : "";
43444504

43454505
try {
43464506
let accessToken = account.accessToken;
@@ -4350,26 +4510,49 @@ while (attempted.size < Math.max(1, accountCount)) {
43504510
typeof account.expiresAt !== "number" ||
43514511
account.expiresAt <= Date.now() + 30_000
43524512
) {
4353-
const refreshResult = await queuedRefresh(account.refreshToken);
4513+
const previousRefreshToken = account.refreshToken;
4514+
if (!previousRefreshToken) {
4515+
throw new Error("Cannot refresh: account has no refresh token");
4516+
}
4517+
const refreshResult = await queuedRefresh(previousRefreshToken);
43544518
if (refreshResult.type !== "success") {
43554519
throw new Error(refreshResult.message ?? refreshResult.reason);
43564520
}
4357-
account.refreshToken = refreshResult.refresh;
4358-
account.accessToken = refreshResult.access;
4359-
account.expiresAt = refreshResult.expires;
4521+
4522+
let refreshedCount = 0;
4523+
for (const storedAccount of storage.accounts) {
4524+
if (!storedAccount) continue;
4525+
if (storedAccount.refreshToken === previousRefreshToken) {
4526+
applyRefreshedCredentials(storedAccount, refreshResult);
4527+
refreshedCount += 1;
4528+
}
4529+
}
4530+
if (refreshedCount === 0) {
4531+
applyRefreshedCredentials(account, refreshResult);
4532+
}
4533+
4534+
const persistedRefresh = await persistRefreshedCredentials({
4535+
previousRefreshToken,
4536+
accountId: account.accountId,
4537+
organizationId: account.organizationId,
4538+
email: account.email,
4539+
refreshResult,
4540+
});
4541+
43604542
accessToken = refreshResult.access;
4361-
storageChanged = true;
4543+
storageChanged = storageChanged || persistedRefresh;
43624544
}
43634545

4364-
const accountId = account.accountId ?? extractAccountId(accessToken);
4546+
const effectiveAccount = sharesActiveCredential ? effectiveDisplayAccount : account;
4547+
const accountId = effectiveAccount.accountId ?? extractAccountId(accessToken);
43654548
if (!accountId) {
43664549
throw new Error("Missing account id");
43674550
}
43684551

43694552
const payload = await fetchUsage({
43704553
accountId,
43714554
accessToken,
4372-
organizationId: account.organizationId,
4555+
organizationId: effectiveAccount.organizationId,
43734556
});
43744557

43754558
const primary = mapWindow(payload.rate_limit?.primary_window ?? null);
@@ -4379,7 +4562,7 @@ while (attempted.size < Math.max(1, accountCount)) {
43794562
payload.additional_rate_limits?.find((entry) => entry.limit_name === "code_review_rate_limit")?.rate_limit ??
43804563
null;
43814564
const codeReview = mapWindow(codeReviewRateLimit?.primary_window ?? null);
4382-
const credits = formatCredits(payload.credits ?? null);
4565+
const credits = formatCredits(payload.credits ?? null);
43834566
const additionalLimits = (payload.additional_rate_limits ?? []).filter(
43844567
(entry) => entry.limit_name !== "code_review_rate_limit",
43854568
);
@@ -4434,12 +4617,7 @@ while (attempted.size < Math.max(1, accountCount)) {
44344617
}
44354618

44364619
if (storageChanged) {
4437-
await saveAccounts(storage);
4438-
if (cachedAccountManager) {
4439-
const reloadedManager = await AccountManager.loadFromDisk();
4440-
cachedAccountManager = reloadedManager;
4441-
accountManagerPromise = Promise.resolve(reloadedManager);
4442-
}
4620+
invalidateAccountManagerCache();
44434621
}
44444622

44454623
while (lines.length > 0 && lines[lines.length - 1] === "") {

0 commit comments

Comments
 (0)