refactor: centralize oauth identity and route refresh workflow

This commit is contained in:
Cita
2026-03-23 11:53:59 +08:00
parent 328b492ab3
commit 3877f3a3f1
22 changed files with 361 additions and 78 deletions
@@ -1,5 +1,6 @@
import { describe, expect, it } from 'vitest';
import {
buildOauthIdentityBackfillPatch,
buildOauthInfoFromAccount,
buildStoredOauthStateFromAccount,
getOauthInfoFromAccount,
@@ -133,4 +134,40 @@ describe('oauth account identity helpers', () => {
quota: { status: 'supported' },
});
});
it('builds a structured identity backfill patch from legacy oauth metadata only for missing columns', () => {
expect(buildOauthIdentityBackfillPatch({
oauthProvider: null,
oauthAccountKey: null,
oauthProjectId: null,
extraConfig: JSON.stringify({
oauth: {
provider: 'codex',
accountKey: 'legacy-account',
projectId: 'legacy-project',
refreshToken: 'refresh-token',
},
}),
})).toEqual({
oauthProvider: 'codex',
oauthAccountKey: 'legacy-account',
oauthProjectId: 'legacy-project',
});
expect(buildOauthIdentityBackfillPatch({
oauthProvider: 'codex',
oauthAccountKey: 'structured-account',
oauthProjectId: null,
extraConfig: JSON.stringify({
oauth: {
provider: 'legacy-provider-ignored',
accountKey: 'legacy-account-ignored',
projectId: 'legacy-project',
refreshToken: 'refresh-token',
},
}),
})).toEqual({
oauthProjectId: 'legacy-project',
});
});
});
+59 -25
View File
@@ -52,6 +52,8 @@ type OauthIdentityCarrier = {
oauthProjectId?: string | null;
};
type StoredOauthIdentity = Pick<OauthInfo, 'provider' | 'accountId' | 'accountKey' | 'projectId'>;
function parseExtraConfig(extraConfig?: string | null): ParsedExtraConfig {
if (!extraConfig) return {};
try {
@@ -113,17 +115,26 @@ function asModelDiscoveryStatus(value: unknown): OauthModelDiscoveryStatus | und
return undefined;
}
function parseStoredOauthState(extraConfig?: string | null): Partial<OauthInfo> | null {
function parseStoredOauthIdentity(extraConfig?: string | null): StoredOauthIdentity | null {
const parsed = parseExtraConfig(extraConfig).oauth;
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return null;
const accountKey = asTrimmedString(parsed.accountKey) || asTrimmedString(parsed.accountId);
const provider = asTrimmedString(parsed.provider);
if (!provider) return null;
return {
provider: asTrimmedString(parsed.provider),
provider,
accountId: asTrimmedString(parsed.accountId) || accountKey,
accountKey,
projectId: asTrimmedString(parsed.projectId),
};
}
function parseStoredOauthRuntimeState(extraConfig?: string | null): Partial<OauthInfo> | null {
const parsed = parseExtraConfig(extraConfig).oauth;
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return null;
return {
email: asTrimmedString(parsed.email),
planType: asTrimmedString(parsed.planType),
projectId: asTrimmedString(parsed.projectId),
tokenExpiresAt: asPositiveInteger(parsed.tokenExpiresAt),
refreshToken: asTrimmedString(parsed.refreshToken),
idToken: asTrimmedString(parsed.idToken),
@@ -137,45 +148,68 @@ function parseStoredOauthState(extraConfig?: string | null): Partial<OauthInfo>
}
export function getOauthInfoFromExtraConfig(extraConfig?: string | null): OauthInfo | null {
const parsed = parseStoredOauthState(extraConfig);
const provider = parsed?.provider;
const identity = parseStoredOauthIdentity(extraConfig);
const runtime = parseStoredOauthRuntimeState(extraConfig);
const provider = identity?.provider;
if (!provider) return null;
return {
provider,
accountId: parsed.accountId || parsed.accountKey,
accountKey: parsed.accountKey,
email: parsed.email,
planType: parsed.planType,
projectId: parsed.projectId,
tokenExpiresAt: parsed.tokenExpiresAt,
refreshToken: parsed.refreshToken,
idToken: parsed.idToken,
providerData: parsed.providerData,
quota: parsed.quota,
modelDiscoveryStatus: parsed.modelDiscoveryStatus,
lastModelSyncAt: parsed.lastModelSyncAt,
lastModelSyncError: parsed.lastModelSyncError,
lastDiscoveredModels: parsed.lastDiscoveredModels,
accountId: identity?.accountId || identity?.accountKey,
accountKey: identity?.accountKey,
projectId: identity?.projectId,
email: runtime?.email,
planType: runtime?.planType,
tokenExpiresAt: runtime?.tokenExpiresAt,
refreshToken: runtime?.refreshToken,
idToken: runtime?.idToken,
providerData: runtime?.providerData,
quota: runtime?.quota,
modelDiscoveryStatus: runtime?.modelDiscoveryStatus,
lastModelSyncAt: runtime?.lastModelSyncAt,
lastModelSyncError: runtime?.lastModelSyncError,
lastDiscoveredModels: runtime?.lastDiscoveredModels,
};
}
export function getOauthInfoFromAccount(account?: OauthIdentityCarrier | null): OauthInfo | null {
if (!account) return null;
const stored = parseStoredOauthState(account.extraConfig);
const provider = asTrimmedString(account.oauthProvider) || stored?.provider;
const storedIdentity = parseStoredOauthIdentity(account.extraConfig);
const storedRuntime = parseStoredOauthRuntimeState(account.extraConfig);
const provider = asTrimmedString(account.oauthProvider) || storedIdentity?.provider;
if (!provider) return null;
const structuredAccountKey = asTrimmedString(account.oauthAccountKey);
const accountKey = structuredAccountKey || stored?.accountKey || stored?.accountId;
const projectId = asTrimmedString(account.oauthProjectId) || stored?.projectId;
const accountKey = structuredAccountKey || storedIdentity?.accountKey || storedIdentity?.accountId;
const projectId = asTrimmedString(account.oauthProjectId) || storedIdentity?.projectId;
return {
...(stored || {}),
...(storedRuntime || {}),
provider,
accountId: structuredAccountKey || stored?.accountId || accountKey,
accountId: structuredAccountKey || storedIdentity?.accountId || accountKey,
accountKey,
projectId,
};
}
export function buildOauthIdentityBackfillPatch(
account?: OauthIdentityCarrier | null,
): Partial<Pick<OauthIdentityCarrier, 'oauthProvider' | 'oauthAccountKey' | 'oauthProjectId'>> | null {
if (!account) return null;
const legacyIdentity = parseStoredOauthIdentity(account.extraConfig);
if (!legacyIdentity?.provider) return null;
const patch: Partial<Pick<OauthIdentityCarrier, 'oauthProvider' | 'oauthAccountKey' | 'oauthProjectId'>> = {};
if (!asTrimmedString(account.oauthProvider)) {
patch.oauthProvider = legacyIdentity.provider;
}
if (!asTrimmedString(account.oauthAccountKey) && (legacyIdentity.accountKey || legacyIdentity.accountId)) {
patch.oauthAccountKey = legacyIdentity.accountKey || legacyIdentity.accountId;
}
if (!asTrimmedString(account.oauthProjectId) && legacyIdentity.projectId) {
patch.oauthProjectId = legacyIdentity.projectId;
}
return Object.keys(patch).length > 0 ? patch : null;
}
export function buildOauthInfo(
extraConfig?: string | null,
patch: Partial<OauthInfo> = {},
@@ -0,0 +1,43 @@
import { eq, isNotNull, isNull, or } from 'drizzle-orm';
import { db, schema } from '../../db/index.js';
import { buildOauthIdentityBackfillPatch } from './oauthAccount.js';
let inFlightOauthIdentityBackfill: Promise<number> | null = null;
async function runOauthIdentityBackfill(): Promise<number> {
const rows = await db.select().from(schema.accounts)
.where(or(
isNotNull(schema.accounts.extraConfig),
isNull(schema.accounts.oauthProvider),
))
.all();
let updated = 0;
for (const row of rows) {
const patch = buildOauthIdentityBackfillPatch(row);
if (!patch) continue;
await db.update(schema.accounts).set({
...patch,
updatedAt: new Date().toISOString(),
}).where(eq(schema.accounts.id, row.id)).run();
updated += 1;
}
return updated;
}
export async function ensureOauthIdentityBackfill(): Promise<number> {
if (inFlightOauthIdentityBackfill) {
return inFlightOauthIdentityBackfill;
}
inFlightOauthIdentityBackfill = (async () => {
try {
return await runOauthIdentityBackfill();
} finally {
inFlightOauthIdentityBackfill = null;
}
})();
return inFlightOauthIdentityBackfill;
}
+7 -4
View File
@@ -1,7 +1,8 @@
import { and, desc, eq, inArray, isNull, or, sql } from 'drizzle-orm';
import { db, schema } from '../../db/index.js';
import { mergeAccountExtraConfig } from '../accountExtraConfig.js';
import { refreshModelsForAccount, rebuildTokenRoutesFromAvailability } from '../modelService.js';
import { refreshModelsForAccount } from '../modelService.js';
import * as routeRefreshWorkflow from '../routeRefreshWorkflow.js';
import {
createOauthSession,
getOauthSession,
@@ -22,6 +23,7 @@ import {
getOauthInfoFromAccount,
} from './oauthAccount.js';
import { buildCodexOauthInfo } from './codexAccount.js';
import { ensureOauthIdentityBackfill } from './oauthIdentityBackfill.js';
import { buildQuotaSnapshotFromOauthInfo, refreshOauthQuotaSnapshot } from './quota.js';
type OAuthProviderMetadata = ReturnType<typeof listOauthProviders>[number];
@@ -404,7 +406,7 @@ export async function handleOauthCallback(input: {
updatedAt: previousAccount.updatedAt,
}).where(eq(schema.accounts.id, previousAccount.id)).run();
}
await rebuildTokenRoutesFromAvailability();
await routeRefreshWorkflow.rebuildRoutesOnly();
const errorMessage = refreshResult.errorMessage || `${input.provider} model discovery failed`;
markOauthSessionError(input.state, errorMessage);
throw new Error(errorMessage);
@@ -417,7 +419,7 @@ export async function handleOauthCallback(input: {
}).where(eq(schema.accounts.id, account.id)).run();
}
await rebuildTokenRoutesFromAvailability();
await routeRefreshWorkflow.rebuildRoutesOnly();
markOauthSessionSuccess(input.state, {
accountId: account.id,
siteId: site.id,
@@ -464,6 +466,7 @@ export async function listOauthConnections(options: {
limit?: number;
offset?: number;
} = {}) {
await ensureOauthIdentityBackfill();
const limit = Math.max(1, Math.min(200, Math.trunc(options.limit ?? 50)));
const offset = Math.max(0, Math.trunc(options.offset ?? 0));
@@ -565,7 +568,7 @@ export async function deleteOauthConnection(accountId: number) {
throw new Error('account is not managed by oauth');
}
await db.delete(schema.accounts).where(eq(schema.accounts.id, accountId)).run();
await rebuildTokenRoutesFromAvailability();
await routeRefreshWorkflow.rebuildRoutesOnly();
return { success: true };
}