smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): preserved fresh codex blocks

Browse Source 
- Delayed healthy-usage reconciliation for newly-set local Codex blocks so lagging /usage responses cannot immediately undo a real 429 backoff.
- Added regression coverage for fresh usage-limit blocks that see healthy usage during selection.
- Kept broker stale-block reconciliation covered by seeding a persisted-only block.

Fixes #4980
This commit is contained in:
roboomp
2026-07-09 21:36:35 +00:00
parent 5b0d670d6b
commit 0ffc43bfc0
2 changed files with 103 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/src/auth-storage.ts
+20
View File
@@ -965,6 +965,8 @@ export class AuthStorage {
#sessionLastCredential: Map<string, Map<string, { type: AuthCredential["type"]; index: number }>> = new Map();
/** Maps provider:type -> credentialIndex -> blockedUntilMs for temporary backoff. */
#credentialBackoff: Map<string, Map<number, number>> = new Map();
/** Earliest time a freshly-set in-memory block may be cleared by live usage reconciliation. */
#credentialBackoffProbeAfter: Map<string, Map<number, number>> = new Map();
#usageProviderResolver?: (provider: Provider) => UsageProvider | undefined;
#rankingStrategyResolver?: (provider: Provider) => CredentialRankingStrategy | undefined;
#usageCache: UsageCache;
@@ -1345,6 +1347,9 @@ export class AuthStorage {
if (backoffMap.size === 0) {
this.#credentialBackoff.delete(backoffKey);
}
const probeAfterMap = this.#credentialBackoffProbeAfter.get(backoffKey);
probeAfterMap?.delete(credentialIndex);
if (probeAfterMap?.size === 0) this.#credentialBackoffProbeAfter.delete(backoffKey);
return undefined;
}
return blockedUntil;
@@ -1437,6 +1442,9 @@ export class AuthStorage {
const nextBlockedUntil = Math.max(existing, blockedUntilMs);
backoffMap.set(credentialIndex, nextBlockedUntil);
this.#credentialBackoff.set(backoffKey, backoffMap);
const probeAfterMap = this.#credentialBackoffProbeAfter.get(backoffKey) ?? new Map<number, number>();
probeAfterMap.set(credentialIndex, Math.min(nextBlockedUntil, Date.now() + USAGE_REPORT_TTL_MS));
this.#credentialBackoffProbeAfter.set(backoffKey, probeAfterMap);
this.#invalidateUsageReportCache(provider);
const upsertCredentialBlock = this.#store.upsertCredentialBlock?.bind(this.#store);
@@ -4385,6 +4393,11 @@ export class AuthStorage {
backoffMap.delete(index);
if (backoffMap.size === 0) this.#credentialBackoff.delete(key);
}
for (const [key, probeAfterMap] of this.#credentialBackoffProbeAfter) {
if (key !== providerKey && !key.startsWith(scopedPrefix)) continue;
probeAfterMap.delete(index);
if (probeAfterMap.size === 0) this.#credentialBackoffProbeAfter.delete(key);
}
}
/**
@@ -4410,6 +4423,13 @@ export class AuthStorage {
const blockScope = this.#rankingStrategyResolver?.(provider)?.blockScope?.({});
const blockedUntilMs = this.#getCredentialBlockedUntil(provider, providerKey, credentialIndex, blockScope);
if (blockedUntilMs === undefined) return;
// `/usage` can lag the request path that just returned 429. Fresh local
// blocks get one usage-cache window before healthy reports may clear them.
const nowMs = Date.now();
const scopedBackoffKey = this.#toScopedBackoffKey(providerKey, blockScope);
const globalProbeAfterMs = this.#credentialBackoffProbeAfter.get(providerKey)?.get(credentialIndex) ?? 0;
const scopedProbeAfterMs = this.#credentialBackoffProbeAfter.get(scopedBackoffKey)?.get(credentialIndex) ?? 0;
if (Math.max(globalProbeAfterMs, scopedProbeAfterMs) > nowMs) return;
this.#clearCredentialBlocks(provider, credentialId);
logger.info("Cleared stale Codex usage-limit block after healthy live usage report", {
credentialId,
packages/ai/test/auth-storage-codex-selection.test.ts
+83 -31
View File
@@ -551,6 +551,73 @@ describe("AuthStorage codex oauth ranking", () => {
expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs);
});
test("keeps a fresh Codex usage-limit block when selection sees healthy usage", async () => {
if (!authStorage || !store?.getCredentialBlock) {
throw new Error("test setup failed");
}
await authStorage.set("openai-codex", [
{ type: "oauth", ...createCredential("acct-fresh-blocked", "fresh-blocked@example.com") },
{ type: "oauth", ...createCredential("acct-fresh-healthy", "fresh-healthy@example.com") },
]);
usageByAccount.set(
"acct-fresh-blocked",
createCodexUsageReport({
accountId: "acct-fresh-blocked",
primary: { usedFraction: 0.2, resetInMs: HOUR_MS },
secondary: { usedFraction: 0.3, resetInMs: WEEK_MS },
metadata: {
allowed: true,
limitReached: false,
planType: "pro",
email: "fresh-blocked@example.com",
accountId: "acct-fresh-blocked",
},
}),
);
usageByAccount.set(
"acct-fresh-healthy",
createCodexUsageReport({
accountId: "acct-fresh-healthy",
primary: { usedFraction: 0.2, resetInMs: HOUR_MS },
secondary: { usedFraction: 0.3, resetInMs: WEEK_MS },
metadata: {
allowed: true,
limitReached: false,
planType: "pro",
email: "fresh-healthy@example.com",
accountId: "acct-fresh-healthy",
},
}),
);
const blockedRow = store.listAuthCredentials("openai-codex").find(row => {
const credential = row.credential;
return credential.type === "oauth" && credential.accountId === "acct-fresh-blocked";
});
if (!blockedRow) throw new Error("expected blocked credential row");
let blockedSessionId: string | undefined;
for (let index = 0; index < 100; index += 1) {
const sessionId = `codex-fresh-block-selected-${index}`;
if ((await authStorage.getApiKey("openai-codex", sessionId)) === "api-acct-fresh-blocked") {
blockedSessionId = sessionId;
break;
}
}
if (!blockedSessionId) throw new Error("expected a session selecting the soon-blocked account");
const markResult = await authStorage.markUsageLimitReached("openai-codex", blockedSessionId, {
retryAfterMs: 6 * 24 * HOUR_MS,
});
expect(markResult.switched).toBe(true);
const selectionAfterBlock = await authStorage.getApiKey("openai-codex", blockedSessionId);
expect(selectionAfterBlock).not.toBe("api-acct-fresh-blocked");
expect(selectionAfterBlock).toBe("api-acct-fresh-healthy");
expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined();
});
test("an older in-flight healthy Codex usage report does not clear a newer usage-limit block", async () => {
if (!authStorage || !store?.getCredentialBlock) {
throw new Error("test setup failed");
@@ -639,7 +706,7 @@ describe("AuthStorage codex oauth ranking", () => {
});
test("broker-sourced healthy Codex usage clears remote gateway backoff", async () => {
if (!authStorage || !store?.getCredentialBlock) {
if (!authStorage || !store?.getCredentialBlock || !store.upsertCredentialBlock) {
throw new Error("test setup failed");
}
@@ -679,6 +746,18 @@ describe("AuthStorage codex oauth ranking", () => {
}),
);
const staleBlockedRow = store.listAuthCredentials("openai-codex").find(row => {
const credential = row.credential;
return credential.type === "oauth" && credential.accountId === "acct-broker-blocked";
});
if (!staleBlockedRow) throw new Error("expected stale blocked credential row");
store.upsertCredentialBlock({
credentialId: staleBlockedRow.id,
providerKey: "openai-codex:oauth",
blockScope: "shared",
blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS,
});
const token = "codex-broker-reconcile";
const handle = startAuthBroker({
storage: authStorage,
@@ -688,18 +767,6 @@ describe("AuthStorage codex oauth ranking", () => {
});
try {
const brokerClient = new AuthBrokerClient({ url: handle.url, token });
const originalUpsertCredentialBlock = brokerClient.upsertCredentialBlock.bind(brokerClient);
const blockPersisted = Promise.withResolvers<void>();
vi.spyOn(brokerClient, "upsertCredentialBlock").mockImplementation(async (id, block, signal) => {
try {
const response = await originalUpsertCredentialBlock(id, block, signal);
blockPersisted.resolve();
return response;
} catch (error) {
blockPersisted.reject(error);
throw error;
}
});
const initialResult = await brokerClient.fetchSnapshot();
if (initialResult.status !== 200) throw new Error("expected broker snapshot");
const blockedRow = initialResult.snapshot.credentials.find(entry => {
@@ -715,31 +782,16 @@ describe("AuthStorage codex oauth ranking", () => {
const clientStorage = new AuthStorage(remoteStore);
await clientStorage.reload();
try {
let blockedSessionId: string | undefined;
for (let index = 0; index < 100; index += 1) {
const sessionId = `broker-codex-local-block-${index}`;
const apiKey = await clientStorage.getApiKey("openai-codex", sessionId);
if (apiKey === "api-acct-broker-blocked") {
blockedSessionId = sessionId;
break;
}
}
if (!blockedSessionId) throw new Error("expected a session selecting the blocked account");
const markResult = await clientStorage.markUsageLimitReached("openai-codex", blockedSessionId, {
retryAfterMs: 6 * 24 * HOUR_MS,
});
expect(markResult.switched).toBe(true);
expect(remoteStore.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined();
await blockPersisted.promise;
expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined();
await clientStorage.fetchUsageReports();
expect(remoteStore.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeUndefined();
expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeUndefined();
expect(await clientStorage.getApiKey("openai-codex", blockedSessionId)).toBe("api-acct-broker-blocked");
expect(await clientStorage.getApiKey("openai-codex", "broker-codex-reconciled")).toBe(
"api-acct-broker-blocked",
);
} finally {
clientStorage.close();
remoteStore.close();