Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix Gemini OAuth refresh and failure coverage

Browse Source
This commit is contained in:
Cita
2026-03-21 18:38:00 +08:00
parent 6a19d809f0
commit dd48550ffa
5 changed files with 634 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/services/modelService.discovery.test.ts
+115
View File
@@ -719,6 +719,121 @@ describe('refreshModelsForAccount credential discovery', () => {
});
});
it('refreshes gemini oauth access token during validation after a 401 and persists the refreshed credentials', async () => {
getApiTokenMock.mockResolvedValue(null);
getModelsMock.mockRejectedValue(new Error('gemini oauth validation should not call adapter.getModels'));
undiciFetchMock
.mockResolvedValueOnce({
ok: false,
status: 401,
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,
json: async () => ({ state: 'ENABLED' }),
text: async () => JSON.stringify({ state: 'ENABLED' }),
});
const site = await db.insert(schema.sites).values({
name: 'gemini-refresh-site',
url: 'https://gemini.example.com',
platform: 'gemini',
status: 'active',
}).returning().get();
const account = await db.insert(schema.accounts).values({
siteId: site.id,
username: 'gemini-user@example.com',
accessToken: 'gemini-access-token-expired',
apiToken: null,
status: 'active',
oauthProvider: 'gemini-cli',
oauthAccountKey: 'gemini-user@example.com',
oauthProjectId: 'project-refresh-demo',
extraConfig: JSON.stringify({
credentialMode: 'session',
oauth: {
provider: 'gemini-cli',
email: 'gemini-user@example.com',
accountId: 'gemini-user@example.com',
accountKey: 'gemini-user@example.com',
projectId: 'project-refresh-demo',
refreshToken: 'gemini-refresh-token',
},
}),
}).returning().get();
const result = await refreshModelsForAccount(account.id);
expect(result).toMatchObject({
accountId: account.id,
refreshed: true,
status: 'success',
modelCount: expect.any(Number),
discoveredByCredential: true,
});
expect(undiciFetchMock).toHaveBeenCalledTimes(5);
expect(String(undiciFetchMock.mock.calls[0]?.[0] || '')).toContain('/projects/project-refresh-demo/services/');
expect(undiciFetchMock.mock.calls[0]?.[1]).toMatchObject({
method: 'GET',
headers: expect.objectContaining({
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({
method: 'GET',
headers: expect.objectContaining({
Authorization: 'Bearer gemini-access-token-refreshed',
}),
});
const latest = await db.select().from(schema.accounts)
.where(eq(schema.accounts.id, account.id))
.get();
expect(latest).toMatchObject({
accessToken: 'gemini-access-token-refreshed',
oauthProvider: 'gemini-cli',
oauthAccountKey: 'gemini-refreshed-user@example.com',
oauthProjectId: 'project-refresh-demo',
});
const parsed = JSON.parse(latest!.extraConfig || '{}');
expect(parsed.oauth).toMatchObject({
provider: 'gemini-cli',
email: 'gemini-refreshed-user@example.com',
refreshToken: 'gemini-refresh-token-next',
projectId: 'project-refresh-demo',
modelDiscoveryStatus: 'healthy',
});
});
it('discovers antigravity oauth models via fetchAvailableModels fallback using the oauth project id', async () => {
getApiTokenMock.mockResolvedValue(null);
getModelsMock.mockRejectedValue(new Error('antigravity oauth discovery should not call adapter.getModels'));
src/server/services/modelService.ts
+97 -8
View File
@@ -24,6 +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 {
ANTIGRAVITY_DAILY_UPSTREAM_BASE_URL,
ANTIGRAVITY_MODELS_USER_AGENT,
@@ -481,6 +482,76 @@ function buildSuccessfulRefreshResult(input: {
};
}
function shouldRetryModelDiscoveryWithOauthRefresh(error: unknown): boolean {
const message = ((error as { message?: string })?.message || '').toLowerCase();
return message.includes('http 401')
|| message.includes('unauthorized')
|| 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))
@@ -661,13 +732,31 @@ export async function refreshModelsForAccount(accountId: number): Promise<ModelR
if (oauth?.provider === 'gemini-cli') {
const checkedAt = new Date().toISOString();
const startedAt = Date.now();
let discoveryAccount = account;
try {
await withTimeout(
() => withAccountProxyOverride(accountProxyUrl,
() => validateGeminiCliOauthConnection({ account })),
MODEL_DISCOVERY_TIMEOUT_MS,
`gemini cli oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`,
);
try {
await withTimeout(
() => withAccountProxyOverride(accountProxyUrl,
() => validateGeminiCliOauthConnection({ account: discoveryAccount })),
MODEL_DISCOVERY_TIMEOUT_MS,
`gemini cli oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`,
);
} catch (error) {
if (!shouldRetryModelDiscoveryWithOauthRefresh(error)) {
throw error;
}
const refreshedAccount = await refreshOauthAccountForModelDiscovery(discoveryAccount);
if (!refreshedAccount) {
throw error;
}
discoveryAccount = refreshedAccount;
await withTimeout(
() => withAccountProxyOverride(accountProxyUrl,
() => validateGeminiCliOauthConnection({ account: discoveryAccount })),
MODEL_DISCOVERY_TIMEOUT_MS,
`gemini cli oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`,
);
}
await db.insert(schema.modelAvailability).values(
GEMINI_CLI_STATIC_MODELS.map((modelName) => ({
accountId,
@@ -678,7 +767,7 @@ export async function refreshModelsForAccount(accountId: number): Promise<ModelR
})),
).run();
await updateOauthModelDiscoveryState({
account,
account: discoveryAccount,
checkedAt,
status: 'healthy',
lastDiscoveredModels: GEMINI_CLI_STATIC_MODELS,
@@ -702,7 +791,7 @@ export async function refreshModelsForAccount(accountId: number): Promise<ModelR
const errorCode = classifyModelDiscoveryError(rawMessage);
const errorMessage = `Gemini CLI 模型获取失败（${rawMessage}）`;
await updateOauthModelDiscoveryState({
account,
account: discoveryAccount,
checkedAt,
status: 'abnormal',
lastModelSyncError: errorMessage,