From ab78d9325213623067b5bd867b35464cd4eefd73 Mon Sep 17 00:00:00 2001 From: vin <49336820+Draculabo@users.noreply.github.com> Date: Wed, 11 Feb 2026 12:36:01 +0800 Subject: [PATCH] fix: project id forwarding regression (#93) * fix(proxy): avoid empty project context and hydrate project_id from account token * test(proxy): add regressions for empty project and token project_id hydration * test(scripts): add live request validator for project_id forwarding --- src/lib/antigravity/ClaudeRequestMapper.ts | 9 +- src/lib/antigravity/types.ts | 2 +- src/server/modules/proxy/proxy.service.ts | 14 +- .../modules/proxy/token-manager.service.ts | 46 ++- .../scripts/project-id-request-validation.ts | 358 ++++++++++++++++++ src/tests/unit/cloud-monitor.test.ts | 96 +++++ src/tests/unit/proxy-retry-mock.test.ts | 22 +- 7 files changed, 534 insertions(+), 13 deletions(-) create mode 100644 src/tests/scripts/project-id-request-validation.ts diff --git a/src/lib/antigravity/ClaudeRequestMapper.ts b/src/lib/antigravity/ClaudeRequestMapper.ts index 2d30ae4..23a1bee 100644 --- a/src/lib/antigravity/ClaudeRequestMapper.ts +++ b/src/lib/antigravity/ClaudeRequestMapper.ts @@ -43,7 +43,7 @@ interface RequestConfig { */ export function transformClaudeRequestIn( claudeReq: ClaudeRequest, - projectId: string, + projectId?: string, ): GeminiInternalRequest { // Check for networking tools (server tool or built-in tool) const hasWebSearchTool = detectsNetworkingTool(claudeReq.tools); @@ -163,8 +163,9 @@ export function transformClaudeRequestIn( const requestId = `agent-${uuidv4()}`; + const normalizedProjectId = projectId?.trim(); + const body: GeminiInternalRequest = { - project: projectId, requestId: requestId, request: innerRequest as GeminiInternalRequest['request'], model: config.finalModel, @@ -172,6 +173,10 @@ export function transformClaudeRequestIn( requestType: config.requestType, }; + if (normalizedProjectId) { + body.project = normalizedProjectId; + } + if (claudeReq.metadata?.user_id) { body.sessionId = claudeReq.metadata.user_id; } diff --git a/src/lib/antigravity/types.ts b/src/lib/antigravity/types.ts index 7efd76c..0c846e0 100644 --- a/src/lib/antigravity/types.ts +++ b/src/lib/antigravity/types.ts @@ -269,7 +269,7 @@ export interface GeminiRequest { } export interface GeminiInternalRequest { - project: string; + project?: string; requestId: string; request: GeminiRequest; model: string; diff --git a/src/server/modules/proxy/proxy.service.ts b/src/server/modules/proxy/proxy.service.ts index 4ab7342..a4cd86c 100644 --- a/src/server/modules/proxy/proxy.service.ts +++ b/src/server/modules/proxy/proxy.service.ts @@ -476,17 +476,24 @@ export class ProxyService { private createGeminiInternalRequest( model: string, request: GeminiRequest, - projectId: string, + projectId: string | undefined, requestType: string, ): GeminiInternalRequest { - return { - project: projectId, + const normalizedProjectId = projectId?.trim(); + + const internalRequest: GeminiInternalRequest = { requestId: uuidv4(), request: this.toInternalGeminiRequest(request), model, userAgent: process.env.PROXY_REQUEST_USER_AGENT ?? this.defaultRequestUserAgent, requestType, }; + + if (normalizedProjectId) { + internalRequest.project = normalizedProjectId; + } + + return internalRequest; } private normalizeGeminiGenerateResponse(response: GeminiResponse): GeminiResponse { @@ -1481,6 +1488,7 @@ export class ProxyService { private isProjectNotFoundError(errorMessage: string): boolean { const msg = errorMessage.toLowerCase(); return ( + msg.includes('invalid project resource name projects/') || (msg.includes('resource projects/') && msg.includes('could not be found')) || (msg.includes('project') && msg.includes('not found')) ); diff --git a/src/server/modules/proxy/token-manager.service.ts b/src/server/modules/proxy/token-manager.service.ts index 3d5c8b5..9aee40e 100644 --- a/src/server/modules/proxy/token-manager.service.ts +++ b/src/server/modules/proxy/token-manager.service.ts @@ -16,6 +16,19 @@ interface TokenData { upstream_proxy_url?: string; } +function normalizeProjectId(projectId: string | null | undefined): string | undefined { + if (typeof projectId !== 'string') { + return undefined; + } + + const trimmedProjectId = projectId.trim(); + if (trimmedProjectId === '' || /^cloud-code-\d+$/i.test(trimmedProjectId)) { + return undefined; + } + + return trimmedProjectId; +} + @Injectable() export class TokenManagerService implements OnModuleInit { private readonly logger = new Logger(TokenManagerService.name); @@ -177,7 +190,7 @@ export class TokenManagerService implements OnModuleInit { tokenData.expiry_timestamp = nowSeconds + newTokens.expires_in; // Save to DB - await this.saveRefreshedToken(accountId, tokenData); + await this.saveTokenState(accountId, tokenData); this.tokens.set(accountId, tokenData); this.logger.log(`Refreshed token for ${tokenData.email}`); @@ -187,13 +200,31 @@ export class TokenManagerService implements OnModuleInit { } if ( - typeof tokenData.project_id !== 'string' || - tokenData.project_id.trim() === '' || - /^cloud-code-\d+$/i.test(tokenData.project_id.trim()) + normalizeProjectId(tokenData.project_id) === undefined ) { tokenData.project_id = undefined; } + if (!tokenData.project_id) { + try { + const fetchedProjectId = await GoogleAPIService.fetchProjectId(tokenData.access_token); + const normalizedProjectId = normalizeProjectId(fetchedProjectId); + + if (normalizedProjectId) { + tokenData.project_id = normalizedProjectId; + await this.saveTokenState(accountId, tokenData); + this.tokens.set(accountId, tokenData); + this.logger.log(`Resolved project ID for ${tokenData.email}: ${normalizedProjectId}`); + } else { + this.logger.warn( + `Project ID is unavailable for ${tokenData.email}; request will continue without project context`, + ); + } + } catch (error) { + this.logger.warn(`Failed to resolve project ID for ${tokenData.email}`, error); + } + } + this.logger.log(`Using account: ${tokenData.email}`); if (sessionKey) { this.sessionBindings.set(sessionKey, { @@ -272,7 +303,7 @@ export class TokenManagerService implements OnModuleInit { ); } - private async saveRefreshedToken(accountId: string, tokenData: TokenData) { + private async saveTokenState(accountId: string, tokenData: TokenData) { try { const acc = await CloudAccountRepo.getAccount(accountId); if (acc && acc.token) { @@ -281,11 +312,14 @@ export class TokenManagerService implements OnModuleInit { access_token: tokenData.access_token, expires_in: tokenData.expires_in, expiry_timestamp: tokenData.expiry_timestamp, + project_id: tokenData.project_id ?? acc.token.project_id, + session_id: tokenData.session_id ?? acc.token.session_id, + upstream_proxy_url: tokenData.upstream_proxy_url ?? acc.token.upstream_proxy_url, }; await CloudAccountRepo.updateToken(accountId, newToken); } } catch (e) { - this.logger.error('Unable to persist refreshed token to DB', e); + this.logger.error('Unable to persist token state to DB', e); } } diff --git a/src/tests/scripts/project-id-request-validation.ts b/src/tests/scripts/project-id-request-validation.ts new file mode 100644 index 0000000..4b87a3d --- /dev/null +++ b/src/tests/scripts/project-id-request-validation.ts @@ -0,0 +1,358 @@ +import assert from 'node:assert/strict'; + +import { transformClaudeRequestIn } from '../../lib/antigravity/ClaudeRequestMapper'; +import { ProxyService } from '../../server/modules/proxy/proxy.service'; +import { TokenManagerService } from '../../server/modules/proxy/token-manager.service'; +import { CloudAccountRepo } from '../../ipc/database/cloudHandler'; +import { GoogleAPIService } from '../../services/GoogleAPIService'; + +const mockTokenManager: any = { + getNextToken: async () => null, + markAsRateLimited: () => undefined, + markAsForbidden: () => undefined, +}; + +const mockGeminiClient: any = { + streamGenerateInternal: async () => undefined, + generateInternal: async () => undefined, +}; + +class TestableProxyService extends ProxyService { + constructor() { + super(mockTokenManager, mockGeminiClient); + } + + public createGeminiInternal( + model: string, + request: Record, + projectId: string | undefined, + requestType: string, + ): Record { + return (this as any).createGeminiInternalRequest(model, request, projectId, requestType); + } +} + +function getEnv(name: string, fallback?: string): string { + const value = process.env[name]; + if (value && value.trim() !== '') { + return value; + } + if (fallback !== undefined) { + return fallback; + } + throw new Error(`Missing required env: ${name}`); +} + +function getEnvFromList(names: string[], fallback?: string): string { + for (const name of names) { + const value = process.env[name]; + if (value && value.trim() !== '') { + return value; + } + } + if (fallback !== undefined) { + return fallback; + } + throw new Error(`Missing required env. Tried: ${names.join(', ')}`); +} + +function getLiveHeaders(apiKey?: string): Record { + const headers: Record = { + 'Content-Type': 'application/json', + }; + if (apiKey && apiKey.trim() !== '') { + headers.Authorization = apiKey.startsWith('Bearer ') ? apiKey : `Bearer ${apiKey}`; + } + return headers; +} + +async function callLive( + baseUrl: string, + path: string, + body: Record, + timeoutMs: number, + apiKey?: string, +): Promise<{ status: number; text: string }> { + const url = `${baseUrl}${path}`; + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), timeoutMs); + + try { + const response = await fetch(url, { + method: 'POST', + headers: getLiveHeaders(apiKey), + body: JSON.stringify(body), + signal: controller.signal, + }); + const text = await response.text(); + return { + status: response.status, + text, + }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + throw new Error( + `Failed to call ${url}. Ensure proxy service is running and reachable. Original error: ${message}`, + ); + } finally { + clearTimeout(timer); + } +} + +async function validateMapperProjectField(): Promise { + const request = { + model: 'claude-sonnet-4-5', + stream: false, + max_tokens: 128, + messages: [{ role: 'user', content: 'hello' }], + }; + + const emptyProjectBody = transformClaudeRequestIn(request as any, ''); + assert.ok(!Object.prototype.hasOwnProperty.call(emptyProjectBody, 'project')); + + const blankProjectBody = transformClaudeRequestIn(request as any, ' '); + assert.ok(!Object.prototype.hasOwnProperty.call(blankProjectBody, 'project')); + + const validProjectBody = transformClaudeRequestIn(request as any, 'resolved-project-id'); + assert.equal(validProjectBody.project, 'resolved-project-id'); +} + +async function validateProxyInternalBuilder(): Promise { + const service = new TestableProxyService(); + + const emptyInternalBody = service.createGeminiInternal( + 'gemini-2.5-flash', + { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] }, + '', + 'generate-content', + ); + console.log( + '[DEBUG] Internal request with empty project_id:', + JSON.stringify(emptyInternalBody, null, 2), + ); + assert.ok(!Object.prototype.hasOwnProperty.call(emptyInternalBody, 'project')); + + const validInternalBody = service.createGeminiInternal( + 'gemini-2.5-flash', + { contents: [{ role: 'user', parts: [{ text: 'hello' }] }] }, + 'resolved-project-id', + 'generate-content', + ); + console.log( + '[DEBUG] Internal request with valid project_id:', + JSON.stringify(validInternalBody, null, 2), + ); + assert.equal(validInternalBody.project, 'resolved-project-id'); +} + +async function validateRuntimeGeminiRequestPath(): Promise { + const service = new TestableProxyService(); + const originalGetNextToken = mockTokenManager.getNextToken; + const originalGenerateInternal = mockGeminiClient.generateInternal; + let capturedBody: Record | null = null; + + mockTokenManager.getNextToken = async () => ({ + id: 'account-1', + email: 'project-test@example.com', + token: { + access_token: 'access-token', + refresh_token: 'refresh-token', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: Math.floor(Date.now() / 1000) + 3600, + project_id: '', + upstream_proxy_url: undefined, + }, + }); + + mockGeminiClient.generateInternal = async (body: Record) => { + capturedBody = body; + return { + candidates: [ + { + content: { + parts: [{ text: 'ok' }], + }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + totalTokenCount: 1, + }, + }; + }; + + try { + await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + } finally { + mockTokenManager.getNextToken = originalGetNextToken; + mockGeminiClient.generateInternal = originalGenerateInternal; + } + + assert.ok(capturedBody, 'Expected captured Gemini internal request'); + console.log( + '[DEBUG] Runtime captured internal request (from handleGeminiGenerateContent):', + JSON.stringify(capturedBody, null, 2), + ); + assert.ok(!Object.prototype.hasOwnProperty.call(capturedBody, 'project')); +} + +async function validateRuntimeGeminiRequestFromAccountProject(): Promise { + const sourceProjectId = `account-project-${Date.now()}`; + const account = { + id: 'account-source-1', + provider: 'google', + email: 'account-source@example.com', + token: { + access_token: 'source-access-token', + refresh_token: 'source-refresh-token', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: Math.floor(Date.now() / 1000) + 3600, + project_id: sourceProjectId, + }, + created_at: 1, + last_used: 1, + }; + + const cloudAccountRepo = CloudAccountRepo as any; + const googleApiService = GoogleAPIService as any; + + const originalGetAccounts = cloudAccountRepo.getAccounts; + const originalGetAccount = cloudAccountRepo.getAccount; + const originalUpdateToken = cloudAccountRepo.updateToken; + const originalFetchProjectId = googleApiService.fetchProjectId; + const originalGenerateInternal = mockGeminiClient.generateInternal; + + let capturedBody: Record | null = null; + let fetchProjectIdCalled = false; + + cloudAccountRepo.getAccounts = async () => [account]; + cloudAccountRepo.getAccount = async (id: string) => (id === account.id ? account : null); + cloudAccountRepo.updateToken = async () => undefined; + googleApiService.fetchProjectId = async () => { + fetchProjectIdCalled = true; + return 'unexpected-project-id'; + }; + + mockGeminiClient.generateInternal = async (body: Record) => { + capturedBody = body; + return { + candidates: [ + { + content: { + parts: [{ text: 'ok' }], + }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + totalTokenCount: 1, + }, + }; + }; + + try { + const tokenManager = new TokenManagerService(); + const service = new ProxyService(tokenManager as any, mockGeminiClient as any); + await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello from account source' }] }], + } as any); + } finally { + cloudAccountRepo.getAccounts = originalGetAccounts; + cloudAccountRepo.getAccount = originalGetAccount; + cloudAccountRepo.updateToken = originalUpdateToken; + googleApiService.fetchProjectId = originalFetchProjectId; + mockGeminiClient.generateInternal = originalGenerateInternal; + } + + assert.ok(capturedBody, 'Expected captured Gemini internal request'); + const captured = capturedBody as Record; + assert.equal( + captured['project'], + sourceProjectId, + 'Expected project field to come from account.token.project_id', + ); + assert.equal( + fetchProjectIdCalled, + false, + 'fetchProjectId should not be called when account already has project_id', + ); + + console.log('[DEBUG] Account source project_id:', sourceProjectId); + console.log( + '[DEBUG] Runtime captured internal request (project from account):', + JSON.stringify(captured, null, 2), + ); +} + +async function validateLiveRequests(): Promise { + const baseUrl = getEnvFromList(['PROJECT_ID_TEST_BASE_URL', 'PARITY_BASE_URL_A']); + const timeoutMs = Number(getEnv('PROJECT_ID_TEST_TIMEOUT_MS', '30000')); + const apiKey = getEnvFromList( + ['PROJECT_ID_TEST_API_KEY', 'PARITY_API_KEY_A', 'PARITY_API_KEY'], + '', + ); + + const gemini = await callLive( + baseUrl, + '/v1beta/models/gemini-2.5-flash:generateContent', + { + contents: [{ role: 'user', parts: [{ text: 'project-id live validation' }] }], + }, + timeoutMs, + apiKey, + ); + assert.ok( + !gemini.text.toLowerCase().includes('invalid project resource name projects/'), + `[gemini-live] still hit invalid empty project path. status=${gemini.status}, body=${gemini.text}`, + ); + console.log(`[PASS] Live Gemini request check (status=${gemini.status})`); + + const anthropic = await callLive( + baseUrl, + '/v1/messages', + { + model: 'claude-sonnet-4-5', + stream: false, + max_tokens: 128, + messages: [{ role: 'user', content: 'project-id live validation' }], + }, + timeoutMs, + apiKey, + ); + assert.ok( + !anthropic.text.toLowerCase().includes('invalid project resource name projects/'), + `[anthropic-live] still hit invalid empty project path. status=${anthropic.status}, body=${anthropic.text}`, + ); + console.log(`[PASS] Live Anthropic request check (status=${anthropic.status})`); +} + +async function main(): Promise { + await validateMapperProjectField(); + console.log('[PASS] ClaudeRequestMapper project field behavior'); + + await validateProxyInternalBuilder(); + console.log('[PASS] ProxyService internal request project field behavior'); + + await validateRuntimeGeminiRequestPath(); + console.log('[PASS] Runtime Gemini request path omits empty project'); + + await validateRuntimeGeminiRequestFromAccountProject(); + console.log('[PASS] Runtime Gemini request uses project_id from account source'); + + if (process.argv.includes('--live')) { + await validateLiveRequests(); + console.log('[PASS] Live request validation'); + } + + console.log('[DONE] Project ID request validation completed'); +} + +main().catch((error) => { + console.error('[FAIL] Project ID request validation failed'); + console.error(error); + process.exit(1); +}); diff --git a/src/tests/unit/cloud-monitor.test.ts b/src/tests/unit/cloud-monitor.test.ts index 3b7d1cf..08701dc 100644 --- a/src/tests/unit/cloud-monitor.test.ts +++ b/src/tests/unit/cloud-monitor.test.ts @@ -3,6 +3,7 @@ import { CloudMonitorService } from '../../services/CloudMonitorService'; import { CloudAccountRepo } from '../../ipc/database/cloudHandler'; import { GoogleAPIService } from '../../services/GoogleAPIService'; import { AutoSwitchService } from '../../services/AutoSwitchService'; +import { TokenManagerService } from '../../server/modules/proxy/token-manager.service'; import { logger } from '../../utils/logger'; // Mock dependencies @@ -192,3 +193,98 @@ describe('CloudMonitorService', () => { }); }); }); + +describe('TokenManagerService project-id hydration', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.useRealTimers(); + }); + + it('hydrates missing project_id and persists it', async () => { + const account = { + id: 'acc-1', + provider: 'google', + email: 'user@example.com', + token: { + access_token: 'access-token', + refresh_token: 'refresh-token', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: Math.floor(Date.now() / 1000) + 3600, + }, + created_at: 1, + last_used: 1, + }; + + vi.mocked(CloudAccountRepo.getAccounts).mockResolvedValue([account] as never); + vi.mocked(CloudAccountRepo.getAccount).mockResolvedValue(account as never); + vi.mocked(CloudAccountRepo.updateToken).mockResolvedValue(undefined as never); + vi.mocked(GoogleAPIService.fetchProjectId).mockResolvedValue('resolved-project' as never); + + const service = new TokenManagerService(); + const selectedToken = await service.getNextToken(); + + expect(GoogleAPIService.fetchProjectId).toHaveBeenCalledWith('access-token'); + expect(CloudAccountRepo.updateToken).toHaveBeenCalledWith( + 'acc-1', + expect.objectContaining({ + project_id: 'resolved-project', + }), + ); + expect(selectedToken?.token.project_id).toBe('resolved-project'); + }); + + it('keeps existing valid project_id without extra fetch', async () => { + const account = { + id: 'acc-2', + provider: 'google', + email: 'existing@example.com', + token: { + access_token: 'access-token-2', + refresh_token: 'refresh-token-2', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: Math.floor(Date.now() / 1000) + 3600, + project_id: 'existing-project', + }, + created_at: 1, + last_used: 1, + }; + + vi.mocked(CloudAccountRepo.getAccounts).mockResolvedValue([account] as never); + vi.mocked(GoogleAPIService.fetchProjectId).mockResolvedValue('new-project' as never); + + const service = new TokenManagerService(); + const selectedToken = await service.getNextToken(); + + expect(GoogleAPIService.fetchProjectId).not.toHaveBeenCalled(); + expect(selectedToken?.token.project_id).toBe('existing-project'); + }); + + it('keeps request usable when project_id cannot be resolved', async () => { + const account = { + id: 'acc-3', + provider: 'google', + email: 'missing@example.com', + token: { + access_token: 'access-token-3', + refresh_token: 'refresh-token-3', + token_type: 'Bearer', + expires_in: 3600, + expiry_timestamp: Math.floor(Date.now() / 1000) + 3600, + project_id: 'cloud-code-123', + }, + created_at: 1, + last_used: 1, + }; + + vi.mocked(CloudAccountRepo.getAccounts).mockResolvedValue([account] as never); + vi.mocked(GoogleAPIService.fetchProjectId).mockResolvedValue(null as never); + + const service = new TokenManagerService(); + const selectedToken = await service.getNextToken(); + + expect(GoogleAPIService.fetchProjectId).toHaveBeenCalledWith('access-token-3'); + expect(selectedToken?.token.project_id).toBeUndefined(); + }); +}); diff --git a/src/tests/unit/proxy-retry-mock.test.ts b/src/tests/unit/proxy-retry-mock.test.ts index 1553321..1475b56 100644 --- a/src/tests/unit/proxy-retry-mock.test.ts +++ b/src/tests/unit/proxy-retry-mock.test.ts @@ -454,10 +454,30 @@ describe('ProxyService Empty Stream Retry Logic', () => { expect(mockTokenManager.getNextToken).toHaveBeenCalledTimes(1); expect(mockGeminiClient.generateInternal).toHaveBeenCalledTimes(2); expect(mockGeminiClient.generateInternal.mock.calls[0][0].project).toBe('project-1'); - expect(mockGeminiClient.generateInternal.mock.calls[1][0].project).toBe(''); + expect(mockGeminiClient.generateInternal.mock.calls[1][0].project).toBeUndefined(); + expect(mockGeminiClient.generateInternal.mock.calls[1][0]).not.toHaveProperty('project'); expect((result as any).candidates?.[0]?.content?.parts?.[0]?.text).toBe('ok'); }); + it('omits empty project id in Gemini internal payload', async () => { + const service = new TestableProxyService(); + const token = createToken('acc-1'); + token.token.project_id = ''; + mockTokenManager.getNextToken.mockResolvedValue(token); + mockGeminiClient.generateInternal.mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'ok' }] }, finishReason: 'STOP' }], + usageMetadata: { totalTokenCount: 5 }, + }); + + await service.handleGeminiGenerateContent('models/gemini-2.5-flash', { + contents: [{ role: 'user', parts: [{ text: 'hello' }] }], + } as any); + + const internalPayload = mockGeminiClient.generateInternal.mock.calls[0][0]; + expect(internalPayload.project).toBeUndefined(); + expect(internalPayload).not.toHaveProperty('project'); + }); + it('uses generate-content requestType for Gemini stream internal payload', async () => { const service = new TestableProxyService(); mockTokenManager.getNextToken.mockResolvedValue(createToken('acc-1'));