Merge runtime token recovery and quota review fixes

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

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