From a02701009773c05615c3430284e3385336fc682e Mon Sep 17 00:00:00 2001 From: Gaurav <39389231+gsquared94@users.noreply.github.com> Date: Fri, 5 Sep 2025 19:55:33 -0700 Subject: [PATCH] fix: handle vpc-sc violations in LoadCodeAssist method (#7824) --- packages/core/src/code_assist/server.test.ts | 37 ++++++++++++++++++ packages/core/src/code_assist/server.ts | 40 +++++++++++++++++--- packages/core/src/code_assist/types.ts | 20 +++++++++- 3 files changed, 90 insertions(+), 7 deletions(-) diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts index bc3978fb14..967493abca 100644 --- a/packages/core/src/code_assist/server.test.ts +++ b/packages/core/src/code_assist/server.test.ts @@ -215,4 +215,41 @@ describe('CodeAssistServer', () => { }), ).rejects.toThrow(); }); + + it('should handle VPC-SC errors when calling loadCodeAssist', async () => { + const client = new OAuth2Client(); + const server = new CodeAssistServer( + client, + 'test-project', + {}, + 'test-session', + UserTierId.FREE, + ); + const mockVpcScError = { + response: { + data: { + error: { + details: [ + { + reason: 'SECURITY_POLICY_VIOLATED', + }, + ], + }, + }, + }, + }; + vi.spyOn(server, 'requestPost').mockRejectedValue(mockVpcScError); + + const response = await server.loadCodeAssist({ + metadata: {}, + }); + + expect(server.requestPost).toHaveBeenCalledWith( + 'loadCodeAssist', + expect.any(Object), + ); + expect(response).toEqual({ + currentTier: { id: UserTierId.STANDARD }, + }); + }); }); diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index cf40fc1aa4..915d07c1df 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -7,6 +7,7 @@ import type { OAuth2Client } from 'google-auth-library'; import type { CodeAssistGlobalUserSettingResponse, + GoogleRpcResponse, LoadCodeAssistRequest, LoadCodeAssistResponse, LongRunningOperationResponse, @@ -23,7 +24,7 @@ import type { } from '@google/genai'; import * as readline from 'node:readline'; import type { ContentGenerator } from '../core/contentGenerator.js'; -import type { UserTierId } from './types.js'; +import { UserTierId } from './types.js'; import type { CaCountTokenResponse, CaGenerateContentResponse, @@ -103,10 +104,20 @@ export class CodeAssistServer implements ContentGenerator { async loadCodeAssist( req: LoadCodeAssistRequest, ): Promise { - return await this.requestPost( - 'loadCodeAssist', - req, - ); + try { + return await this.requestPost( + 'loadCodeAssist', + req, + ); + } catch (e) { + if (isVpcScAffectedUser(e)) { + return { + currentTier: { id: UserTierId.STANDARD }, + }; + } else { + throw e; + } + } } async getCodeAssistGlobalUserSetting(): Promise { @@ -221,3 +232,22 @@ export class CodeAssistServer implements ContentGenerator { return `${endpoint}/${CODE_ASSIST_API_VERSION}:${method}`; } } + +function isVpcScAffectedUser(error: unknown): boolean { + if (error && typeof error === 'object' && 'response' in error) { + const gaxiosError = error as { + response?: { + data?: unknown; + }; + }; + const response = gaxiosError.response?.data as + | GoogleRpcResponse + | undefined; + if (Array.isArray(response?.error?.details)) { + return response.error.details.some( + (detail) => detail.reason === 'SECURITY_POLICY_VIOLATED', + ); + } + } + return false; +} diff --git a/packages/core/src/code_assist/types.ts b/packages/core/src/code_assist/types.ts index 367186857b..b79094bbb8 100644 --- a/packages/core/src/code_assist/types.ts +++ b/packages/core/src/code_assist/types.ts @@ -57,8 +57,8 @@ export interface LoadCodeAssistResponse { */ export interface GeminiUserTier { id: UserTierId; - name: string; - description: string; + name?: string; + description?: string; // This value is used to declare whether a given tier requires the user to configure the project setting on the IDE settings or not. userDefinedCloudaicompanionProject?: boolean | null; isDefault?: boolean; @@ -183,3 +183,19 @@ export interface CodeAssistGlobalUserSettingResponse { cloudaicompanionProject?: string; freeTierDataCollectionOptin: boolean; } + +/** + * Relevant fields that can be returned from a Google RPC response + */ +export interface GoogleRpcResponse { + error?: { + details?: GoogleRpcErrorInfo[]; + }; +} + +/** + * Relevant fields that can be returned in the details of an error returned from GoogleRPCs + */ +interface GoogleRpcErrorInfo { + reason?: string; +}