From 554d80c060ea079c7edf99ff58aae46466521472 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 11 May 2026 10:46:53 +0100 Subject: [PATCH] fix(codex): scope auth state updates --- .../codex/src/app-server/run-attempt.test.ts | 37 +++++++++- .../codex/src/app-server/run-attempt.ts | 69 +++++++++++++++---- .../run.overflow-compaction.harness.ts | 5 +- .../run.overflow-compaction.test.ts | 7 ++ src/agents/pi-embedded-runner/run.ts | 19 ++++- 5 files changed, 119 insertions(+), 18 deletions(-) diff --git a/extensions/codex/src/app-server/run-attempt.test.ts b/extensions/codex/src/app-server/run-attempt.test.ts index 5c45919c8e5e..c6fc6bfade46 100644 --- a/extensions/codex/src/app-server/run-attempt.test.ts +++ b/extensions/codex/src/app-server/run-attempt.test.ts @@ -2162,6 +2162,7 @@ describe("runCodexAppServerAttempt", () => { const sessionFile = path.join(tempDir, "session.jsonl"); const workspaceDir = path.join(tempDir, "workspace"); const resetsAt = Math.ceil(Date.now() / 1000) + 120; + const authProfileId = "openai-codex:work"; const harnessRef: { current?: ReturnType } = {}; const harness = createStartedThreadHarness(async (method) => { if (method === "turn/start") { @@ -2177,7 +2178,22 @@ describe("runCodexAppServerAttempt", () => { }); harnessRef.current = harness; - const result = await runCodexAppServerAttempt(createParams(sessionFile, workspaceDir)); + const params = createParams(sessionFile, workspaceDir); + params.authProfileId = authProfileId; + params.authProfileStore = { + version: 1, + profiles: { + [authProfileId]: { + type: "oauth", + provider: "openai-codex", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }, + }; + + const result = await runCodexAppServerAttempt(params); expect(result.promptErrorSource).toBe("prompt"); expect(result.promptError).toContain("You've reached your Codex subscription usage limit."); expect(result.promptError).toContain("Next reset in"); @@ -2187,6 +2203,7 @@ describe("runCodexAppServerAttempt", () => { const sessionFile = path.join(tempDir, "session.jsonl"); const workspaceDir = path.join(tempDir, "workspace"); const resetsAt = Math.ceil(Date.now() / 1000) + 120; + const authProfileId = "openai-codex:work"; rememberCodexRateLimits({ rateLimits: { limitId: "codex", @@ -2208,13 +2225,29 @@ describe("runCodexAppServerAttempt", () => { return undefined; }); - const run = runCodexAppServerAttempt(createParams(sessionFile, workspaceDir)); + const params = createParams(sessionFile, workspaceDir); + params.authProfileId = authProfileId; + params.authProfileStore = { + version: 1, + profiles: { + [authProfileId]: { + type: "oauth", + provider: "openai-codex", + access: "access", + refresh: "refresh", + expires: Date.now() + 60_000, + }, + }, + }; + + const run = runCodexAppServerAttempt(params); await harness.waitForMethod("turn/start"); const result = await run; expect(result.promptErrorSource).toBe("prompt"); expect(result.promptError).toContain("You've reached your Codex subscription usage limit."); expect(result.promptError).toContain("Next reset in"); + expect(params.authProfileStore.usageStats?.[authProfileId]?.blockedUntil).toBeUndefined(); }); it("refreshes Codex account rate limits when turn/start omits reset details", async () => { diff --git a/extensions/codex/src/app-server/run-attempt.ts b/extensions/codex/src/app-server/run-attempt.ts index 5bc91ea32cc5..a18b8753e852 100644 --- a/extensions/codex/src/app-server/run-attempt.ts +++ b/extensions/codex/src/app-server/run-attempt.ts @@ -1368,7 +1368,7 @@ export async function runCodexAppServerAttempt( timeoutMs: appServer.requestTimeoutMs, signal: runAbortController.signal, }); - const turnStartErrorMessage = usageLimitError ?? formatErrorMessage(error); + const turnStartErrorMessage = usageLimitError?.message ?? formatErrorMessage(error); emitCodexAppServerEvent(params, { stream: "codex_app_server.lifecycle", data: { phase: "turn_start_failed", error: turnStartErrorMessage }, @@ -1419,13 +1419,14 @@ export async function runCodexAppServerAttempt( }); params.abortSignal?.removeEventListener("abort", abortFromUpstream); if (usageLimitError) { - await markCodexAuthProfileBlockedFromRecentRateLimits({ + await markCodexAuthProfileBlockedFromRateLimits({ params, authProfileId: startupAuthProfileId, + rateLimits: usageLimitError.rateLimitsForProfile, }); return buildCodexTurnStartFailureResult({ params, - message: usageLimitError, + message: usageLimitError.message, messagesSnapshot: turnStartFailureMessages, systemPromptReport, }); @@ -1679,15 +1680,16 @@ export async function runCodexAppServerAttempt( } } -async function markCodexAuthProfileBlockedFromRecentRateLimits(params: { +async function markCodexAuthProfileBlockedFromRateLimits(params: { params: EmbeddedRunAttemptParams; authProfileId?: string; + rateLimits?: JsonValue; }): Promise { const authProfileId = params.authProfileId?.trim(); if (!authProfileId || !params.params.authProfileStore) { return; } - const blockedUntil = resolveCodexUsageLimitResetAtMs(readRecentCodexRateLimits()); + const blockedUntil = resolveCodexUsageLimitResetAtMs(params.rateLimits); if (!blockedUntil) { return; } @@ -2212,6 +2214,12 @@ type CodexUsageLimitErrorSource = { message?: string | null; codexErrorInfo?: JsonValue | null; rateLimits?: JsonValue; + rateLimitsTrustedForProfile?: boolean; +}; + +type CodexUsageLimitErrorResult = { + message: string; + rateLimitsForProfile?: JsonValue; }; async function formatCodexTurnStartUsageLimitError(params: { @@ -2220,8 +2228,8 @@ async function formatCodexTurnStartUsageLimitError(params: { pendingNotifications: CodexServerNotification[]; timeoutMs?: number; signal?: AbortSignal; -}): Promise { - return refreshCodexUsageLimitErrorMessage({ +}): Promise { + return refreshCodexUsageLimitError({ client: params.client, source: readCodexTurnStartUsageLimitErrorSource(params.error, params.pendingNotifications), timeoutMs: params.timeoutMs, @@ -2235,9 +2243,32 @@ async function refreshCodexUsageLimitErrorMessage(params: { timeoutMs?: number; signal?: AbortSignal; }): Promise { + return ( + await refreshCodexUsageLimitError({ + client: params.client, + source: params.source, + timeoutMs: params.timeoutMs, + signal: params.signal, + }) + )?.message; +} + +async function refreshCodexUsageLimitError(params: { + client: CodexAppServerClient; + source: CodexUsageLimitErrorSource; + timeoutMs?: number; + signal?: AbortSignal; +}): Promise { const initialMessage = formatCodexUsageLimitErrorMessage(params.source); if (!shouldRefreshCodexRateLimitsForUsageLimitMessage(initialMessage)) { - return initialMessage ?? undefined; + return initialMessage + ? { + message: initialMessage, + ...(params.source.rateLimitsTrustedForProfile + ? { rateLimitsForProfile: params.source.rateLimits } + : {}), + } + : undefined; } const rateLimits = await readCodexRateLimitsFromAppServerForUsageLimitError({ client: params.client, @@ -2245,14 +2276,22 @@ async function refreshCodexUsageLimitErrorMessage(params: { signal: params.signal, }); if (!rateLimits) { - return initialMessage; + return initialMessage + ? { + message: initialMessage, + ...(params.source.rateLimitsTrustedForProfile + ? { rateLimitsForProfile: params.source.rateLimits } + : {}), + } + : undefined; } const refreshedMessage = formatCodexUsageLimitErrorMessage({ message: params.source.message, codexErrorInfo: params.source.codexErrorInfo, rateLimits, }); - return refreshedMessage ?? initialMessage; + const message = refreshedMessage ?? initialMessage; + return message ? { message, rateLimitsForProfile: rateLimits } : undefined; } async function readCodexRateLimitsFromAppServerForUsageLimitError(params: { @@ -2290,14 +2329,16 @@ function readCodexTurnStartUsageLimitErrorSource( pendingNotifications: CodexServerNotification[], ): CodexUsageLimitErrorSource { const notificationError = readLatestCodexErrorNotification(pendingNotifications); + const notificationRateLimits = readLatestRateLimitNotificationPayload(pendingNotifications); const errorPayload = readCodexErrorPayload(error); + const rateLimits = + notificationRateLimits ?? errorPayload.rateLimits ?? readRecentCodexRateLimits(); return { message: notificationError?.message ?? errorPayload.message ?? formatErrorMessage(error), codexErrorInfo: notificationError?.codexErrorInfo ?? errorPayload.codexErrorInfo, - rateLimits: - readLatestRateLimitNotificationPayload(pendingNotifications) ?? - errorPayload.rateLimits ?? - readRecentCodexRateLimits(), + rateLimits, + rateLimitsTrustedForProfile: + notificationRateLimits !== undefined || errorPayload.rateLimits !== undefined, }; } diff --git a/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts b/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts index 89c5e14305a6..7a8f61468a62 100644 --- a/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts +++ b/src/agents/pi-embedded-runner/run.overflow-compaction.harness.ts @@ -223,6 +223,7 @@ export const mockedGetApiKeyForModel = vi.fn( export const mockedEnsureAuthProfileStore = vi.fn(() => ({})); export const mockedEnsureAuthProfileStoreWithoutExternalProfiles = vi.fn(() => ({})); export const mockedResolveAuthProfileOrder = vi.fn(() => [] as string[]); +export const mockedMarkAuthProfileSuccess = vi.fn(async () => {}); export const mockedShouldPreferExplicitConfigApiKeyAuth = vi.fn(() => false); export const overflowBaseRunParams = { @@ -407,6 +408,8 @@ export function resetRunOverflowCompactionHarnessMocks(): void { mockedEnsureAuthProfileStoreWithoutExternalProfiles.mockReturnValue({}); mockedResolveAuthProfileOrder.mockReset(); mockedResolveAuthProfileOrder.mockReturnValue([]); + mockedMarkAuthProfileSuccess.mockReset(); + mockedMarkAuthProfileSuccess.mockResolvedValue(undefined); mockedShouldPreferExplicitConfigApiKeyAuth.mockReset(); mockedShouldPreferExplicitConfigApiKeyAuth.mockReturnValue(false); mockedRunPostCompactionSideEffects.mockReset(); @@ -455,7 +458,7 @@ export async function loadRunOverflowCompactionHarness(): Promise<{ vi.doMock("../auth-profiles.js", () => ({ isProfileInCooldown: vi.fn(() => false), markAuthProfileFailure: vi.fn(async () => {}), - markAuthProfileSuccess: vi.fn(async () => {}), + markAuthProfileSuccess: mockedMarkAuthProfileSuccess, resolveProfilesUnavailableReason: vi.fn(() => undefined), })); diff --git a/src/agents/pi-embedded-runner/run.overflow-compaction.test.ts b/src/agents/pi-embedded-runner/run.overflow-compaction.test.ts index 1cb61c383b0b..b4a5ea0639db 100644 --- a/src/agents/pi-embedded-runner/run.overflow-compaction.test.ts +++ b/src/agents/pi-embedded-runner/run.overflow-compaction.test.ts @@ -22,6 +22,7 @@ import { mockedEnsureAuthProfileStoreWithoutExternalProfiles, mockedGlobalHookRunner, mockedGetApiKeyForModel, + mockedMarkAuthProfileSuccess, mockedPickFallbackThinkingLevel, mockedResolveAuthProfileOrder, mockedResolveContextWindowInfo, @@ -462,6 +463,12 @@ describe("runEmbeddedPiAgent overflow compaction trigger routing", () => { }); const harnessParams = pluginRunAttempt.mock.calls[0]?.[0]; expect(harnessParams?.runtimePlan).toBe(runtimePlan); + expect(mockedMarkAuthProfileSuccess).toHaveBeenCalledWith( + expect.objectContaining({ + provider: "openai-codex", + profileId: "openai-codex:work", + }), + ); }); it("keeps auto-selected OpenAI Codex auth profiles for forced codex harness runs", async () => { diff --git a/src/agents/pi-embedded-runner/run.ts b/src/agents/pi-embedded-runner/run.ts index 1c8d6cc7974c..f3fd34e3ce62 100644 --- a/src/agents/pi-embedded-runner/run.ts +++ b/src/agents/pi-embedded-runner/run.ts @@ -2957,7 +2957,11 @@ export async function runEmbeddedPiAgent( if (lastProfileId) { await markAuthProfileSuccess({ store: profileFailureStore, - provider, + provider: resolveAuthProfileStateProvider( + profileFailureStore, + lastProfileId, + provider, + ), profileId: lastProfileId, agentDir: params.agentDir, }); @@ -3104,3 +3108,16 @@ export async function runEmbeddedPiAgent( }); }); } + +function resolveAuthProfileStateProvider( + store: AuthProfileStore, + profileId: string, + fallbackProvider: string, +): string { + const profileProvider = store.profiles?.[profileId]?.provider?.trim(); + if (profileProvider) { + return profileProvider; + } + const idProvider = profileId.split(":", 1)[0]?.trim(); + return idProvider || fallbackProvider; +}