Merge runtime token recovery and quota review fixes

Co-Authored-By: Paperclip <noreply@paperclip.ing>

# Conflicts:
#	ui/src/components/ProviderQuotaCard.tsx
#	ui/src/pages/Costs.production.tsx
#	ui/src/pages/Costs.tsx

doc/connections/AI-CONNECTIONS.md

@@ -546,7 +546,12 @@ key and fsynced under `<instance-root>/quota-credential-recovery/<company-id>/`
 before vault, grant, or activity writes. Storage and encryption are checked before
 exchanging a single-use refresh token. Failed saves retry without another OAuth
 exchange; persistent failures keep the encrypted record for the next quota poll,
-including after a restart. The journal is removed only after database commit (or
+including after a restart. New OpenAI subscription runtimes serialize credential
+reads with quota exchanges and recover a matching pending replacement before
+materializing an auth home. Authentication-failure handling also recovers a
+matching replacement before marking the grant invalid. A failed recovery save
+defers these actions and leaves both the active grant and journal intact.
+The journal is removed only after database commit (or
 when an authorized newer credential makes it obsolete). Replay checks the grant,
 connection, secret, and original credential fingerprint and cannot reactivate a
 revoked grant or overwrite a reconnect. Back up this directory and the instance

server/src/__tests__/ai-connections.test.ts

@@ -143,6 +143,47 @@ describe("managed AI connections", () => {
     }
   });
 
+  it.each(["runtime", "failed_run"] as const)("recovers pending quota credentials before %s can reuse or invalidate the consumed token", async (consumer) => {
+    const owner = `quota-consumer-${consumer}`;
+    await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
+    const auth = (marker: string) => JSON.stringify({ tokens: { access_token: `${marker}-access`, refresh_token: `${marker}-refresh`, account_id: owner } });
+    await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth("old"));
+    const [row] = await service.quotaAccounts(companyId, owner);
+    const prepare = () => prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local",
+      binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
+    const oldRun = await prepare();
+    const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential")!;
+    const journal = quotaCredentialRecovery(companyId, row.grant.id);
+    await journal.write({ companyId, grantId: row.grant.id, connectionId: row.connection.id, secretId: ref.secretId,
+      baseHash: quotaCredentialHash(auth("old")), value: auth("new") });
+    const fail = () => service.markAuthenticationFailed({ companyId, agentId, runStartedAt: new Date(),
+      attribution: { ...oldRun.attribution, identity: oldRun.identity } });
+    await db.execute(sql.raw(`create function reject_consumer_recovery() returns trigger language plpgsql as $$ begin
+      if new.secret_id = '${ref.secretId}'::uuid then raise exception 'recovery unavailable'; end if; return new; end $$;
+      create trigger reject_consumer_recovery before insert on company_secret_versions for each row execute function reject_consumer_recovery();`));
+    try {
+      await expect(consumer === "runtime" ? prepare() : fail()).rejects.toThrow();
+      expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
+      expect(await service.credential(row)).toBe(auth("old"));
+      expect(await journal.read()).not.toBeNull();
+    } finally {
+      await db.execute(sql`drop trigger reject_consumer_recovery on company_secret_versions; drop function reject_consumer_recovery()`);
+    }
+    const request = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("Recovery must not exchange again"));
+    let nextRun: Awaited<ReturnType<typeof prepare>> | undefined;
+    try {
+      if (consumer === "failed_run") await fail();
+      nextRun = await prepare();
+      expect(await readFile(path.join(String(nextRun.config.env.CODEX_HOME), "auth.json"), "utf8")).toBe(auth("new"));
+      await fail(); // A late failure of the old run cannot invalidate the recovered identity.
+      expect((await db.select().from(connectionGrants).where(eq(connectionGrants.id, row.grant.id)))[0].status).toBe("active");
+      expect(await service.credential(row)).toBe(auth("new"));
+      expect(await journal.read()).toBeNull();
+      expect(request).not.toHaveBeenCalled();
+      expect(await db.select().from(companySecretVersions).where(eq(companySecretVersions.secretId, ref.secretId))).toHaveLength(2);
+    } finally { request.mockRestore(); await oldRun.cleanup(); await nextRun?.cleanup(); await journal.clear(); }
+  });
+
   it.each(["rollback", "lost_commit_reply"] as const)("automatically saves one exchanged token across %s", async (failure) => {
     const owner = `quota-once-${failure}`;
     await db.insert(companyMemberships).values({ companyId, principalId: owner, principalType: "user", status: "active", membershipRole: "member" });
@@ -235,6 +276,9 @@ describe("managed AI connections", () => {
     await db.insert(companyMemberships).values([quotaOwner, runOwner].map(principalId => ({ companyId, principalId, principalType: "user" as const, status: "active" as const, membershipRole: "member" as const })));
     for (const owner of [quotaOwner, runOwner]) await service.save(companyId, owner, { provider: "openai", method: "subscription", ownership: "personal", name: owner, loginSessionId: "fixture", agentIds: [], allAgents: true }, auth(owner, "old"));
     const [quotaRow] = await service.quotaAccounts(companyId, quotaOwner);
+    // Runtime credential reads now serialize with quota exchanges. Prepare this
+    // already-running, different account before testing its contended write-back.
+    const preparedRun = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: runOwner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
     const started = deferredSignal(), release = deferredSignal();
     const request = vi.spyOn(globalThis, "fetch").mockImplementation(async () => {
       started.resolve(); await release.promise;
@@ -246,8 +290,7 @@ describe("managed AI connections", () => {
     let lock: ReturnType<typeof vi.spyOn> | undefined;
     try {
       await started.promise;
-      // This different account starts after the refresh's active-run check.
-      run = await prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: runOwner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
+      run = preparedRun;
       const authFile = path.join(String(run.config.env.CODEX_HOME), "auth.json");
       await writeFile(authFile, auth(runOwner, "replacement"));
       // Inject the two 30-second acquisition expirations, then use the real
@@ -303,24 +346,26 @@ describe("managed AI connections", () => {
       return Response.json({ access_token: "new-access", refresh_token: "new-refresh" });
     });
     let runtime: Awaited<ReturnType<typeof prepareManagedAiRuntime>> | undefined;
+    let lock: ReturnType<typeof vi.spyOn> | undefined;
     const refreshing = service.refreshQuotaCredential(row, original, new AbortController().signal);
     try {
       await exchangeStarted.promise;
       const edit = secretService(db).update(secretId, { description: "Changed during refresh" });
+      lock = vi.spyOn(codexAdapter, "withAccountHomeSecretMutationLock");
       const preparing = prepareManagedAiRuntime(db, { companyId, agentId, responsibleUserId: owner, adapterType: "codex_local", binding: { provider: "openai", method: "subscription", mode: "responsible_user" }, config: {} });
-      // Wait for the runtime's old-token read to reach its bookkeeping update.
-      // After the refresh commits, resolution must detect the changed version.
-      await vi.waitFor(async () => {
-        const waiting = await db.execute(sql`select 1 from pg_stat_activity where datname = current_database() and wait_event_type = 'Lock' and query like '%last_resolved_at%'`);
-        expect(waiting.length).toBeGreaterThan(0);
-      });
+      // Runtime reads now wait at the company file lock before touching the
+      // old credential. The returned home must contain the committed rotation.
+      let prepared = false;
+      void preparing.then(() => { prepared = true; });
+      await vi.waitFor(() => expect(lock).toHaveBeenCalled());
+      expect(prepared).toBe(false);
       releaseExchange.resolve();
       const results = await Promise.all([refreshing, edit, preparing]);
       runtime = results[2];
       expect(JSON.parse(await readFile(path.join(String(runtime.config.env.CODEX_HOME), "auth.json"), "utf8"))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } });
       expect(JSON.parse(await service.credential(row))).toMatchObject({ tokens: { access_token: "new-access", refresh_token: "new-refresh" } });
       expect((await secretService(db).getById(secretId))?.description).toBe("Changed during refresh");
-    } finally { releaseExchange.resolve(); await runtime?.cleanup(); request.mockRestore(); }
+    } finally { releaseExchange.resolve(); await runtime?.cleanup(); request.mockRestore(); lock?.mockRestore(); }
   }, 10_000);
 
   it("serializes quota credential refresh, persists rotation, and defers while the provider is active", async () => {

server/src/services/ai-connection-runtime.ts

@@ -252,7 +252,7 @@ export async function prepareManagedAiRuntime(
       throw unprocessable(
         "The selected default changed. Retry this execution.",
       );
-    const value = await service.credential(selection);
+    const value = await service.runtimeCredential(selection);
     home = await mkdtemp(
       path.join(
         os.tmpdir(),

server/src/services/ai-connections.ts

@@ -816,6 +816,54 @@ export function aiConnectionService(db: Db) {
       return { connectionId: id, grantId };
     }));
   }
+  async function recoverQuotaCredentialInTransaction(tx: Db, row: Parameters<typeof credential>[0], current: string) {
+    const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai);
+    if (!metadata.success || metadata.data.provider !== "openai" || metadata.data.method !== "subscription") return null;
+    const recovery = quotaCredentialRecovery(row.connection.companyId, row.grant.id);
+    const pending = await recovery.read();
+    if (!pending) return null;
+    const ref = row.grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
+    if (!ref || pending.secretId !== ref.secretId || pending.connectionId !== row.connection.id || pending.baseHash !== quotaCredentialHash(current)) {
+      await recovery.clear();
+      return null;
+    }
+    await secretService(tx).rotate(ref.secretId, { value: pending.value }, { userId: row.grant.subjectUserId });
+    await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, row.grant.id));
+    await logActivity(tx, { companyId: row.connection.companyId, actorType: "system", actorId: "quota_recovery",
+      action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id,
+      details: { provider: "openai", grantId: row.grant.id } });
+    return pending.value;
+  }
+
+  /** Never hand a consumed token to a new provider while its replacement awaits commit. */
+  async function runtimeCredential(row: Parameters<typeof credential>[0]) {
+    const metadata = aiConnectionMetadataSchema.safeParse(row.connection.config.ai);
+    if (!metadata.success || metadata.data.provider !== "openai" || metadata.data.method !== "subscription") return credential(row);
+    return withAccountHomeSecretMutationLock(undefined, row.connection.companyId, async () => {
+      let recovered = false;
+      const value = await db.transaction(async tx => {
+        const [grant] = await tx.select().from(connectionGrants).where(and(
+          eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, row.connection.companyId),
+          eq(connectionGrants.connectionId, row.connection.id),
+        )).for("update", { noWait: true });
+        if (!grant || grant.status !== "active") throw new Error("credentials_unavailable");
+        const ref = grant.credentialSecretRefs.find(ref => ref.configPath === "ai.credential");
+        if (!ref) throw new Error("credentials_unavailable");
+        await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
+          eq(companySecrets.id, ref.secretId), eq(companySecrets.companyId, row.connection.companyId),
+        )).for("update", { noWait: true });
+        const selected = { connection: row.connection, grant };
+        const current = await aiConnectionService(tx as unknown as Db).credential(selected);
+        const pending = await recoverQuotaCredentialInTransaction(tx as unknown as Db, selected, current);
+        recovered = pending !== null;
+        return pending ?? current;
+      });
+      if (recovered) await quotaCredentialRecovery(row.connection.companyId, row.grant.id).clear()
+        .catch(error => logger.warn({ err: error, grantId: row.grant.id }, "Saved runtime credential recovery cleanup pending"));
+      return value;
+    });
+  }
+
   /** A late failure must never invalidate credentials that were refreshed or reconnected meanwhile. */
   async function markAuthenticationFailed(input: {
     companyId: string;
@@ -824,41 +872,50 @@ export function aiConnectionService(db: Db) {
     runStartedAt: Date;
     attribution: AiConnectionAttribution & { identity: string };
   }) {
-    return withAccountHomeSecretMutationLock(undefined, input.companyId, () => db.transaction(async (tx) => {
-      const { attribution } = input;
-      const [grant] = await tx.select().from(connectionGrants).where(and(
-        eq(connectionGrants.companyId, input.companyId),
-        eq(connectionGrants.id, attribution.grantId),
-        eq(connectionGrants.connectionId, attribution.connectionId),
-      )).for("update");
-      if (!grant || grant.status !== "active" || grant.updatedAt > input.runStartedAt) return;
-      const [connection] = await tx.select().from(toolConnections).where(and(
-        eq(toolConnections.companyId, input.companyId),
-        eq(toolConnections.id, attribution.connectionId),
-        eq(toolConnections.connectionPurpose, "ai"),
-      ));
-      if (!connection) return;
-      const metadata = aiConnectionMetadataSchema.safeParse(connection.config.ai);
-      if (!metadata.success || metadata.data.provider !== attribution.provider || metadata.data.method !== attribution.method) return;
-      const ref = grant.credentialSecretRefs.find((candidate) => candidate.configPath === "ai.credential");
-      if (!ref) return;
-      await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
-        eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, ref.secretId),
-      )).for("update");
-      const value = await aiConnectionService(tx as unknown as Db).credential({ connection, grant });
-      const generation = createHash("sha256").update(value).digest("hex").slice(0, 16);
-      if (attribution.identity !== `${grant.id}:${attribution.responsibleUserId ?? "shared"}:${generation}`) return;
-      await tx.update(connectionGrants).set({ status: "needs_reauthorization", updatedAt: new Date() })
-        .where(eq(connectionGrants.id, grant.id));
-      await tx.update(toolConnections).set({ healthStatus: "error", healthMessage: "Sign in again to restore this AI connection.", updatedAt: new Date() })
-        .where(eq(toolConnections.id, connection.id));
-      await logActivity(tx as unknown as Db, {
-        companyId: input.companyId, actorType: "system", actorId: input.runId ? "heartbeat" : "adapter_test",
-        agentId: input.agentId, runId: input.runId, action: "ai_connection.authentication_failed",
-        entityType: "tool_connection", entityId: connection.id,
-        details: { provider: attribution.provider, grantId: grant.id },
+    return withAccountHomeSecretMutationLock(undefined, input.companyId, async () => {
+      let recovered = false;
+      await db.transaction(async (tx) => {
+        const { attribution } = input;
+        const [grant] = await tx.select().from(connectionGrants).where(and(
+          eq(connectionGrants.companyId, input.companyId),
+          eq(connectionGrants.id, attribution.grantId),
+          eq(connectionGrants.connectionId, attribution.connectionId),
+        )).for("update");
+        if (!grant || grant.status !== "active" || grant.updatedAt > input.runStartedAt) return;
+        const [connection] = await tx.select().from(toolConnections).where(and(
+          eq(toolConnections.companyId, input.companyId),
+          eq(toolConnections.id, attribution.connectionId),
+          eq(toolConnections.connectionPurpose, "ai"),
+        ));
+        if (!connection) return;
+        const metadata = aiConnectionMetadataSchema.safeParse(connection.config.ai);
+        if (!metadata.success || metadata.data.provider !== attribution.provider || metadata.data.method !== attribution.method) return;
+        const ref = grant.credentialSecretRefs.find((candidate) => candidate.configPath === "ai.credential");
+        if (!ref) return;
+        await tx.select({ id: companySecrets.id }).from(companySecrets).where(and(
+          eq(companySecrets.companyId, input.companyId), eq(companySecrets.id, ref.secretId),
+        )).for("update");
+        const value = await aiConnectionService(tx as unknown as Db).credential({ connection, grant });
+        const generation = createHash("sha256").update(value).digest("hex").slice(0, 16);
+        if (attribution.identity !== `${grant.id}:${attribution.responsibleUserId ?? "shared"}:${generation}`) return;
+        // A quota exchange may have consumed this run's token before a failed
+        // database commit. Save its working replacement instead of revoking it.
+        recovered = await recoverQuotaCredentialInTransaction(tx as unknown as Db, { connection, grant }, value) !== null;
+        if (recovered) return;
+        await tx.update(connectionGrants).set({ status: "needs_reauthorization", updatedAt: new Date() })
+          .where(eq(connectionGrants.id, grant.id));
+        await tx.update(toolConnections).set({ healthStatus: "error", healthMessage: "Sign in again to restore this AI connection.", updatedAt: new Date() })
+          .where(eq(toolConnections.id, connection.id));
+        await logActivity(tx as unknown as Db, {
+          companyId: input.companyId, actorType: "system", actorId: input.runId ? "heartbeat" : "adapter_test",
+          agentId: input.agentId, runId: input.runId, action: "ai_connection.authentication_failed",
+          entityType: "tool_connection", entityId: connection.id,
+          details: { provider: attribution.provider, grantId: grant.id },
+        });
       });
-    }));
+      if (recovered) await quotaCredentialRecovery(input.companyId, input.attribution.grantId).clear()
+        .catch(error => logger.warn({ err: error, grantId: input.attribution.grantId }, "Saved failed-run credential recovery cleanup pending"));
+    });
   }
   /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */
   async function refreshQuotaCredential(row: Parameters<typeof credential>[0], failedValue: string, signal: AbortSignal) {
@@ -968,5 +1025,5 @@ export function aiConnectionService(db: Db) {
       return [{ ...row, summary }];
     });
   }
-  return { list, quotaAccounts, refreshQuotaCredential, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed };
+  return { list, quotaAccounts, refreshQuotaCredential, select, credential, runtimeCredential, probeUsage, save, setDefault, membership, markAuthenticationFailed };
 }

ui/src/components/AccountQuotaPanels.test.tsx

@@ -1,5 +1,6 @@
 import { renderToStaticMarkup } from "react-dom/server";
 import { describe, expect, it } from "vitest";
+import { formatDateTime } from "../lib/utils";
 import { AccountQuotaPanels } from "./AccountQuotaPanels";
 
 describe("account quota panels", () => {
@@ -18,7 +19,7 @@ describe("account quota panels", () => {
       { provider: "openai", accountKey: "a", ok: true, windows: [{ label: "Credits", usedPercent: null, resetsAt: "2026-10-05T00:00:00Z", valueLabel: "$5 remaining" }] },
     ]} />);
     expect(html).toContain("$5 remaining");
-    expect(html).toContain(`Resets ${new Date("2026-10-05T00:00:00Z").toLocaleString()}`);
+    expect(html).toContain(`Resets ${formatDateTime("2026-10-05T00:00:00Z")}`);
     expect(html).not.toContain("0%");
     expect(html).not.toContain('role="progressbar"');
   });

ui/src/components/AccountQuotaPanels.tsx

@@ -1,4 +1,5 @@
 import type { ProviderQuotaResult } from "@paperclipai/shared";
+import { formatDateTime } from "../lib/utils";
 import { QuotaBar } from "./QuotaBar";
 import { quotaUnavailableMessage } from "../lib/quota-refresh";
 
@@ -32,11 +33,12 @@ export function AccountQuotaPanels({
                 <span>{window.label}</span>
                 <span className="text-muted-foreground">{window.valueLabel ?? "Usage not reported"}</span>
               </div>
-              {window.resetsAt && <p className="text-xs text-muted-foreground">Resets {new Date(window.resetsAt).toLocaleString()}</p>}
+              {window.resetsAt && <p className="text-xs font-mono text-muted-foreground">Resets {formatDateTime(window.resetsAt)}</p>}
             </div>
           ) : (
             <QuotaBar
               key={`${window.label}:${index}`}
+              className="font-mono"
               label={window.label}
               percentUsed={window.usedPercent}
               leftLabel={
@@ -47,14 +49,14 @@ export function AccountQuotaPanels({
               }
               rightLabel={
                 window.resetsAt
-                  ? `Resets ${new Date(window.resetsAt).toLocaleString()}`
+                  ? `Resets ${formatDateTime(window.resetsAt)}`
                   : undefined
               }
             />
           ))}
           {account.capturedAt && (
-            <p className="text-xs text-muted-foreground">
-              Last checked {new Date(account.capturedAt).toLocaleString()}
+            <p className="text-xs font-mono text-muted-foreground">
+              Last checked {formatDateTime(account.capturedAt)}
             </p>
           )}
         </section>

ui/src/pages/Costs.quota.test.tsx

@@ -34,6 +34,20 @@ describe("Costs account quota polling", () => {
   });
   afterEach(() => { act(() => root.unmount()); client.clear(); container.remove(); vi.unstubAllGlobals(); });
   const flush = async () => { await act(async () => { await new Promise(resolve => setTimeout(resolve, 10)); }); };
+  it("keeps quota readings with a safe notice when the endpoint rejects", async () => {
+    await act(async () => root.render(<MemoryRouter><SidebarProvider><QueryClientProvider client={client}><Costs /></QueryClientProvider></SidebarProvider></MemoryRouter>));
+    await flush();
+    const tab = [...container.querySelectorAll<HTMLButtonElement>('[role="tab"]')].find(button => button.textContent === "Providers")!;
+    await act(async () => { tab.dispatchEvent(new MouseEvent("mousedown", { bubbles: true, button: 0 })); tab.click(); });
+    await flush();
+    expect(container.textContent).toContain("37% used");
+    api.quotaWindows.mockRejectedValue(new Error("Raw provider diagnostic: secret command"));
+    await act(async () => { await client.invalidateQueries({ queryKey }); }); await flush();
+    expect(container.textContent).toContain("37% used");
+    expect(container.textContent).toContain("Showing the last available quota");
+    expect(container.textContent).not.toContain("Raw provider diagnostic");
+    expect(client.getQueryData(queryKey)).toMatchObject([{ capturedAt: observed.capturedAt }]);
+  });
   it.each(["provider_unavailable", "permission_denied"])("retains the observation on %s but clears rejected or rotated credentials", async (errorFamily) => {
     await act(async () => root.render(<MemoryRouter><SidebarProvider><QueryClientProvider client={client}><Costs initialTab="providers" /></QueryClientProvider></SidebarProvider></MemoryRouter>));
     await flush();