fix: recover managed quota credentials and preserve account displays

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

server/src/services/ai-connections.ts

@@ -1,6 +1,6 @@
 import { syncConnectionCredentialBindings } from "./connection-credential-bindings.js";
 import { createHash, randomUUID } from "node:crypto";
-import { and, eq, inArray, or } from "drizzle-orm";
+import { and, eq, inArray, or, sql } from "drizzle-orm";
 import {
   type Db,
   authUsers,
@@ -8,6 +8,7 @@ import {
   aiConnectionDefaults,
   aiProviderDefaults,
   agents,
+  heartbeatRuns,
   companyMemberships,
   companySecrets,
   userSecretDefinitions,
@@ -848,6 +849,63 @@ export function aiConnectionService(db: Db) {
       });
     });
   }
+  /** Refresh only the selected vaulted identity, serializing rotation and reconnect. */
+  async function refreshQuotaCredential(row: Parameters<typeof credential>[0], failedValue: string, signal: AbortSignal) {
+    return db.transaction(async (tx) => {
+      const companyId = row.connection.companyId;
+      const [grant] = await tx.select().from(connectionGrants).where(and(
+        eq(connectionGrants.id, row.grant.id), eq(connectionGrants.companyId, companyId),
+        eq(connectionGrants.connectionId, row.connection.id),
+      )).for("update");
+      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, companyId),
+      )).for("update");
+      const current = await aiConnectionService(tx as unknown as Db).credential({ connection: row.connection, grant });
+      signal.throwIfAborted();
+      // Another poll, an agent, or reconnect already supplied a newer credential.
+      if (current !== failedValue) return current;
+      // An active provider may rotate this single-use token itself. Defer polling
+      // until its credential write-back instead of invalidating its live copy.
+      const [active] = await tx.select({ id: heartbeatRuns.id }).from(heartbeatRuns)
+        .innerJoin(agents, and(eq(agents.id, heartbeatRuns.agentId), eq(agents.companyId, companyId))).where(and(
+        eq(heartbeatRuns.companyId, companyId), inArray(heartbeatRuns.status, ["queued", "running"]),
+        or(sql`${heartbeatRuns.contextSnapshot}->'aiConnection'->>'grantId' = ${grant.id}`,
+          sql`${agents.runtimeConfig}->'aiConnection'->>'provider' = 'openai'`, eq(agents.adapterType, "codex_local")),
+      )).limit(1);
+      if (active) throw new Error("provider_unavailable");
+      const auth = JSON.parse(current);
+      const refreshToken = auth.tokens?.refresh_token;
+      if (typeof refreshToken !== "string" || !refreshToken) throw new Error("authentication_required");
+      // Matches the Codex OAuth client (openai/codex, login/src/auth/manager.rs).
+      const response = await fetch("https://auth.openai.com/oauth/token", {
+        method: "POST", redirect: "error", signal,
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify({ client_id: "app_EMoamEEZ73f0CkXaXp7hrann", grant_type: "refresh_token", refresh_token: refreshToken }),
+      });
+      if (!response.ok) {
+        await response.body?.cancel();
+        throw new Error(response.status === 400 || response.status === 401 ? "authentication_required" : "provider_unavailable");
+      }
+      const tokens = await response.json() as Record<string, unknown>;
+      if (typeof tokens.access_token !== "string" || !tokens.access_token) throw new Error("provider_unavailable");
+      const refreshed = JSON.stringify({ ...auth, tokens: { ...auth.tokens,
+        access_token: tokens.access_token,
+        ...(typeof tokens.refresh_token === "string" && tokens.refresh_token ? { refresh_token: tokens.refresh_token } : {}),
+        ...(typeof tokens.id_token === "string" && tokens.id_token ? { id_token: tokens.id_token } : {}),
+      }, last_refresh: new Date().toISOString() });
+      // Persist the rotated token even if the quota deadline expires after the
+      // successful exchange; losing it would strand all subsequent executions.
+      await secretService(tx).rotate(ref.secretId, { value: refreshed }, { userId: grant.subjectUserId });
+      await tx.update(connectionGrants).set({ updatedAt: new Date() }).where(eq(connectionGrants.id, grant.id));
+      await logActivity(tx as unknown as Db, { companyId, actorType: "system", actorId: "quota",
+        action: "ai_connection.credential_refreshed", entityType: "tool_connection", entityId: row.connection.id,
+        details: { provider: "openai", grantId: grant.id } });
+      return refreshed;
+    });
+  }
   // A quota read uses the same credential audience as execution. Operator status
   // alone never grants access to another member's personal subscription.
   async function quotaAccounts(companyId: string, userId: string) {
@@ -860,5 +918,5 @@ export function aiConnectionService(db: Db) {
       return [{ ...row, summary }];
     });
   }
-  return { list, quotaAccounts, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed };
+  return { list, quotaAccounts, refreshQuotaCredential, select, credential, probeUsage, save, setDefault, membership, markAuthenticationFailed };
 }

server/src/services/quota-windows.ts

@@ -35,7 +35,8 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId
     const base = { provider: row.summary.provider, accountKey, accountLabel: row.summary.name, source: "managed-connection" };
     const definitionUnavailable = row.grant.credentialSecretRefs.some(ref => {
       const revision = revisions.find(value => value.id === ref.secretId);
-      return revision?.definitionId && (revision.definitionStatus !== "active" || revision.definitionDeletedAt !== null);
+      return revision && (revision.status !== "active" || revision.versionStatus === "disabled" || revision.revokedAt != null
+        || (revision.definitionId && (revision.definitionStatus !== "active" || revision.definitionDeletedAt !== null)));
     });
     if (row.summary.status !== "connected" || definitionUnavailable) {
       return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] };
@@ -52,23 +53,34 @@ export async function fetchCompanyQuotaWindows(db: Db, companyId: string, userId
           credentialResolved = true;
           let windows;
           if (base.provider === "openai") {
-            const auth = JSON.parse(value);
+            let auth;
+            try { auth = JSON.parse(value); } catch { throw new Error("credentials_unavailable"); }
             const token = auth.tokens?.access_token ?? auth.accessToken;
-            const accountId = auth.tokens?.account_id ?? auth.accountId;
-            if (typeof token !== "string" || !token || typeof accountId !== "string" || !accountId) {
+            if (typeof token !== "string" || !token) {
               return { ...base, ok: false, errorFamily: "credentials_unavailable", error: publicError, windows: [] };
             }
-            windows = await fetchCodexQuota(token, accountId, controller.signal);
+            const read = (auth: { tokens?: { access_token?: string; account_id?: string | null }; accessToken?: string; accountId?: string | null }) =>
+              fetchCodexQuota(auth.tokens?.access_token ?? auth.accessToken!, auth.tokens?.account_id ?? auth.accountId ?? null, controller.signal);
+            try { windows = await read(auth); }
+            catch (error) {
+              if (!(error instanceof Error) || !/\b401\b/.test(error.message)) throw error;
+              auth = JSON.parse(await service.refreshQuotaCredential(row, value, controller.signal));
+              controller.signal.throwIfAborted();
+              windows = await read(auth);
+            }
           } else {
             windows = await fetchClaudeQuota(value, controller.signal);
           }
           return { ...base, ok: true, windows, capturedAt: new Date().toISOString() };
         } catch (error) {
           const message = error instanceof Error ? error.message : "Quota unavailable";
-          const invalid = !credentialResolved || /401|403|refresh_token_(?:reused|expired|invalidated)/i.test(message);
+          const missing = message === "credentials_unavailable";
+          const invalid = message === "authentication_required" || (credentialResolved && /\b401\b|refresh_token_(?:reused|expired|invalidated)/i.test(message));
+          const errorFamily = missing ? "credentials_unavailable" : invalid ? "authentication_required"
+            : credentialResolved && /\b403\b/.test(message) ? "permission_denied" : "provider_unavailable";
           logger.warn({ companyId, accountKey, provider: base.provider, authenticationFailed: invalid }, "Connected account quota unavailable");
           return { ...base, ok: false, windows: [], error: publicError,
-            ...(invalid ? { errorFamily: "authentication_required" } : {}) };
+            errorFamily };
         }
       })();
       const bounded = withQuotaTimeout(base.provider, result, () => controller.abort()).then(quota => ({ ...base, ...quota, ...(quota.ok ? {} : { error: publicError }) }));