fix(codex): scope auth state updates

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<typeof createStartedThreadHarness> } = {};
     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 () => {

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<void> {
   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<string | undefined> {
-  return refreshCodexUsageLimitErrorMessage({
+}): Promise<CodexUsageLimitErrorResult | undefined> {
+  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<string | undefined> {
+  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<CodexUsageLimitErrorResult | undefined> {
   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,
   };
 }
 

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),
   }));
 

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 () => {

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;
+}