forked from boos/BoosAPI
fix: reuse oauth singleflight in gemini discovery
This commit is contained in:
Cita
@@ -8,6 +8,7 @@ const getApiTokenMock = vi.fn();
|
||||
const getModelsMock = vi.fn();
|
||||
const undiciFetchMock = vi.fn();
|
||||
const proxyAgentCtorMock = vi.fn();
|
||||
const refreshOauthAccessTokenSingleflightMock = vi.fn();
|
||||
|
||||
class MockProxyAgent {
|
||||
readonly proxyUrl: string;
|
||||
@@ -33,6 +34,10 @@ vi.mock('undici', () => ({
|
||||
Agent: MockAgent,
|
||||
}));
|
||||
|
||||
vi.mock('./oauth/refreshSingleflight.js', () => ({
|
||||
refreshOauthAccessTokenSingleflight: (...args: unknown[]) => refreshOauthAccessTokenSingleflightMock(...args),
|
||||
}));
|
||||
|
||||
type DbModule = typeof import('../db/index.js');
|
||||
type ModelServiceModule = typeof import('./modelService.js');
|
||||
|
||||
@@ -62,6 +67,7 @@ describe('refreshModelsForAccount credential discovery', () => {
|
||||
getModelsMock.mockReset();
|
||||
undiciFetchMock.mockReset();
|
||||
proxyAgentCtorMock.mockReset();
|
||||
refreshOauthAccessTokenSingleflightMock.mockReset();
|
||||
|
||||
await db.delete(schema.routeChannels).run();
|
||||
await db.delete(schema.tokenRoutes).run();
|
||||
@@ -719,9 +725,39 @@ describe('refreshModelsForAccount credential discovery', () => {
|
||||
});
|
||||
});
|
||||
|
||||
it('refreshes gemini oauth access token during validation after a 401 and persists the refreshed credentials', async () => {
|
||||
it('refreshes gemini oauth access token during validation through singleflight and reuses the refreshed account state', async () => {
|
||||
getApiTokenMock.mockResolvedValue(null);
|
||||
getModelsMock.mockRejectedValue(new Error('gemini oauth validation should not call adapter.getModels'));
|
||||
refreshOauthAccessTokenSingleflightMock.mockImplementation(async (accountId: number) => {
|
||||
const refreshedExtraConfig = JSON.stringify({
|
||||
credentialMode: 'session',
|
||||
oauth: {
|
||||
provider: 'gemini-cli',
|
||||
email: 'gemini-refreshed-user@example.com',
|
||||
accountId: 'gemini-refreshed-user@example.com',
|
||||
accountKey: 'gemini-refreshed-user@example.com',
|
||||
projectId: 'project-refresh-demo',
|
||||
refreshToken: 'gemini-refresh-token-next',
|
||||
},
|
||||
});
|
||||
|
||||
await db.update(schema.accounts).set({
|
||||
accessToken: 'gemini-access-token-refreshed',
|
||||
oauthProvider: 'gemini-cli',
|
||||
oauthAccountKey: 'gemini-refreshed-user@example.com',
|
||||
oauthProjectId: 'project-refresh-demo',
|
||||
extraConfig: refreshedExtraConfig,
|
||||
status: 'active',
|
||||
updatedAt: '2026-03-21T00:00:00.000Z',
|
||||
}).where(eq(schema.accounts.id, accountId)).run();
|
||||
|
||||
return {
|
||||
accountId,
|
||||
accessToken: 'gemini-access-token-refreshed',
|
||||
accountKey: 'gemini-refreshed-user@example.com',
|
||||
extraConfig: refreshedExtraConfig,
|
||||
};
|
||||
});
|
||||
undiciFetchMock
|
||||
.mockResolvedValueOnce({
|
||||
ok: false,
|
||||
@@ -729,30 +765,6 @@ describe('refreshModelsForAccount credential discovery', () => {
|
||||
json: async () => ({ error: 'expired' }),
|
||||
text: async () => 'expired',
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({
|
||||
access_token: 'gemini-access-token-refreshed',
|
||||
refresh_token: 'gemini-refresh-token-next',
|
||||
expires_in: 3600,
|
||||
token_type: 'Bearer',
|
||||
scope: 'cloud-platform',
|
||||
}),
|
||||
text: async () => JSON.stringify({ ok: true }),
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({ state: 'ENABLED' }),
|
||||
text: async () => JSON.stringify({ state: 'ENABLED' }),
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
status: 200,
|
||||
json: async () => ({ email: 'gemini-refreshed-user@example.com' }),
|
||||
text: async () => JSON.stringify({ email: 'gemini-refreshed-user@example.com' }),
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
ok: true,
|
||||
status: 200,
|
||||
@@ -798,7 +810,8 @@ describe('refreshModelsForAccount credential discovery', () => {
|
||||
modelCount: expect.any(Number),
|
||||
discoveredByCredential: true,
|
||||
});
|
||||
expect(undiciFetchMock).toHaveBeenCalledTimes(5);
|
||||
expect(refreshOauthAccessTokenSingleflightMock).toHaveBeenCalledWith(account.id);
|
||||
expect(undiciFetchMock).toHaveBeenCalledTimes(2);
|
||||
expect(String(undiciFetchMock.mock.calls[0]?.[0] || '')).toContain('/projects/project-refresh-demo/services/');
|
||||
expect(undiciFetchMock.mock.calls[0]?.[1]).toMatchObject({
|
||||
method: 'GET',
|
||||
@@ -806,9 +819,11 @@ describe('refreshModelsForAccount credential discovery', () => {
|
||||
Authorization: 'Bearer gemini-access-token-expired',
|
||||
}),
|
||||
});
|
||||
expect(String(undiciFetchMock.mock.calls[1]?.[0] || '')).toContain('oauth2.googleapis.com/token');
|
||||
expect(String(undiciFetchMock.mock.calls[4]?.[0] || '')).toContain('/projects/project-refresh-demo/services/');
|
||||
expect(undiciFetchMock.mock.calls[4]?.[1]).toMatchObject({
|
||||
expect(
|
||||
undiciFetchMock.mock.calls.some(([url]) => String(url || '').includes('oauth2.googleapis.com/token')),
|
||||
).toBe(false);
|
||||
expect(String(undiciFetchMock.mock.calls[1]?.[0] || '')).toContain('/projects/project-refresh-demo/services/');
|
||||
expect(undiciFetchMock.mock.calls[1]?.[1]).toMatchObject({
|
||||
method: 'GET',
|
||||
headers: expect.objectContaining({
|
||||
Authorization: 'Bearer gemini-access-token-refreshed',
|
||||
|
||||
@@ -24,7 +24,7 @@ import { withAccountProxyOverride, withExplicitProxyRequestInit, withSiteRecordP
|
||||
import { getCodexOauthInfoFromExtraConfig, isCodexPlatform } from './oauth/codexAccount.js';
|
||||
import { buildOauthInfo, getOauthInfoFromExtraConfig } from './oauth/oauthAccount.js';
|
||||
import { CLAUDE_DEFAULT_ANTHROPIC_VERSION } from './oauth/claudeProvider.js';
|
||||
import { getOAuthProviderDefinition } from './oauth/providers.js';
|
||||
import { refreshOauthAccessTokenSingleflight } from './oauth/refreshSingleflight.js';
|
||||
import {
|
||||
ANTIGRAVITY_DAILY_UPSTREAM_BASE_URL,
|
||||
ANTIGRAVITY_MODELS_USER_AGENT,
|
||||
@@ -489,69 +489,6 @@ function shouldRetryModelDiscoveryWithOauthRefresh(error: unknown): boolean {
|
||||
|| message.includes('unauthenticated');
|
||||
}
|
||||
|
||||
async function refreshOauthAccountForModelDiscovery(
|
||||
account: typeof schema.accounts.$inferSelect,
|
||||
): Promise<typeof schema.accounts.$inferSelect | null> {
|
||||
const oauth = getOauthInfoFromExtraConfig(account.extraConfig);
|
||||
if (!oauth?.provider || !oauth.refreshToken) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const definition = getOAuthProviderDefinition(oauth.provider);
|
||||
if (!definition) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const refreshed = await definition.refreshAccessToken({
|
||||
refreshToken: oauth.refreshToken,
|
||||
oauth: {
|
||||
projectId: oauth.projectId,
|
||||
providerData: oauth.providerData,
|
||||
},
|
||||
});
|
||||
const nextOauth = buildOauthInfo(account.extraConfig, {
|
||||
provider: oauth.provider,
|
||||
accountId: refreshed.accountId || oauth.accountId,
|
||||
accountKey: refreshed.accountKey || oauth.accountKey || refreshed.accountId || oauth.accountId,
|
||||
email: refreshed.email || oauth.email,
|
||||
planType: refreshed.planType || oauth.planType,
|
||||
projectId: refreshed.projectId || oauth.projectId,
|
||||
refreshToken: refreshed.refreshToken || oauth.refreshToken,
|
||||
tokenExpiresAt: refreshed.tokenExpiresAt || oauth.tokenExpiresAt,
|
||||
idToken: refreshed.idToken || oauth.idToken,
|
||||
providerData: {
|
||||
...(oauth.providerData || {}),
|
||||
...(refreshed.providerData || {}),
|
||||
},
|
||||
});
|
||||
const extraConfig = mergeAccountExtraConfig(account.extraConfig, {
|
||||
credentialMode: 'session',
|
||||
oauth: nextOauth,
|
||||
});
|
||||
const updatedAt = new Date().toISOString();
|
||||
|
||||
await db.update(schema.accounts).set({
|
||||
accessToken: refreshed.accessToken,
|
||||
oauthProvider: oauth.provider,
|
||||
oauthAccountKey: nextOauth.accountKey || nextOauth.accountId || null,
|
||||
oauthProjectId: nextOauth.projectId || null,
|
||||
extraConfig,
|
||||
status: 'active',
|
||||
updatedAt,
|
||||
}).where(eq(schema.accounts.id, account.id)).run();
|
||||
|
||||
return {
|
||||
...account,
|
||||
accessToken: refreshed.accessToken,
|
||||
oauthProvider: oauth.provider,
|
||||
oauthAccountKey: nextOauth.accountKey || nextOauth.accountId || null,
|
||||
oauthProjectId: nextOauth.projectId || null,
|
||||
extraConfig,
|
||||
status: 'active',
|
||||
updatedAt,
|
||||
};
|
||||
}
|
||||
|
||||
export async function refreshModelsForAccount(accountId: number): Promise<ModelRefreshResult> {
|
||||
const row = await db.select().from(schema.accounts)
|
||||
.innerJoin(schema.sites, eq(schema.accounts.siteId, schema.sites.id))
|
||||
@@ -745,11 +682,15 @@ export async function refreshModelsForAccount(accountId: number): Promise<ModelR
|
||||
if (!shouldRetryModelDiscoveryWithOauthRefresh(error)) {
|
||||
throw error;
|
||||
}
|
||||
const refreshedAccount = await refreshOauthAccountForModelDiscovery(discoveryAccount);
|
||||
if (!refreshedAccount) {
|
||||
const refreshed = await refreshOauthAccessTokenSingleflight(discoveryAccount.id);
|
||||
if (!refreshed?.extraConfig) {
|
||||
throw error;
|
||||
}
|
||||
discoveryAccount = refreshedAccount;
|
||||
discoveryAccount = {
|
||||
...discoveryAccount,
|
||||
accessToken: refreshed.accessToken,
|
||||
extraConfig: refreshed.extraConfig,
|
||||
};
|
||||
await withTimeout(
|
||||
() => withAccountProxyOverride(accountProxyUrl,
|
||||
() => validateGeminiCliOauthConnection({ account: discoveryAccount })),
|
||||
|
||||
Reference in New Issue
Block a user