fix: handle vpc-sc violations in LoadCodeAssist method (google-gemini#7824)

Author: gsquared94

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 },
+    });
+  });
 });

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<LoadCodeAssistResponse> {
-    return await this.requestPost<LoadCodeAssistResponse>(
-      'loadCodeAssist',
-      req,
-    );
+    try {
+      return await this.requestPost<LoadCodeAssistResponse>(
+        'loadCodeAssist',
+        req,
+      );
+    } catch (e) {
+      if (isVpcScAffectedUser(e)) {
+        return {
+          currentTier: { id: UserTierId.STANDARD },
+        };
+      } else {
+        throw e;
+      }
+    }
   }
 
   async getCodeAssistGlobalUserSetting(): Promise<CodeAssistGlobalUserSettingResponse> {
@@ -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;
+}

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;
+}