Merge "Login with Google Workspace" auth option into "Login with Google" (#1574)

Co-authored-by: Scott Densmore <[email protected]>

12 files changed, 18 insertions, 94 deletions

diff --git a/packages/cli/src/config/auth.ts b/packages/cli/src/config/auth.ts
index 51178513..df47596f 100644
--- a/packages/cli/src/config/auth.ts
+++ b/packages/cli/src/config/auth.ts
@@ -13,13 +13,6 @@ export const validateAuthMethod = (authMethod: string): string | null => {
     return null;
   }
 
-  if (authMethod === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE) {
-    if (!process.env.GOOGLE_CLOUD_PROJECT) {
-      return 'GOOGLE_CLOUD_PROJECT environment variable not found. Add that to your .env and try again, no reload needed!';
-    }
-    return null;
-  }
-
   if (authMethod === AuthType.USE_GEMINI) {
     if (!process.env.GEMINI_API_KEY) {
       return 'GEMINI_API_KEY environment variable not found. Add that to your .env and try again, no reload needed!';
diff --git a/packages/cli/src/ui/components/AuthDialog.tsx b/packages/cli/src/ui/components/AuthDialog.tsx
index 747a83bd..f8ddfaed 100644
--- a/packages/cli/src/ui/components/AuthDialog.tsx
+++ b/packages/cli/src/ui/components/AuthDialog.tsx
@@ -29,32 +29,15 @@ export function AuthDialog({
   const [errorMessage, setErrorMessage] = useState<string | null>(
     initialErrorMessage || null,
   );
-  const allAuthItems = [
+  const items = [
     {
       label: 'Login with Google',
       value: AuthType.LOGIN_WITH_GOOGLE_PERSONAL,
     },
     { label: 'Gemini API Key', value: AuthType.USE_GEMINI },
-    {
-      label: 'Login with Google (for Workspace or licensed Code Assist users)',
-      value: AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE,
-    },
     { label: 'Vertex AI', value: AuthType.USE_VERTEX_AI },
   ];
 
-  const isSelectedAuthInMore = allAuthItems
-    .slice(2)
-    .some((item) => item.value === settings.merged.selectedAuthType);
-
-  const [showAll, setShowAll] = useState(isSelectedAuthInMore);
-
-  const initialAuthItems = [
-    ...allAuthItems.slice(0, 2),
-    { label: 'More...', value: 'more' },
-  ];
-
-  const items = showAll ? allAuthItems : initialAuthItems;
-
   let initialAuthIndex = items.findIndex(
     (item) => item.value === settings.merged.selectedAuthType,
   );
@@ -64,10 +47,6 @@ export function AuthDialog({
   }
 
   const handleAuthSelect = (authMethod: string) => {
-    if (authMethod === 'more') {
-      setShowAll(true);
-      return;
-    }
     const error = validateAuthMethod(authMethod);
     if (error) {
       setErrorMessage(error);
diff --git a/packages/cli/src/ui/hooks/useAuthCommand.ts b/packages/cli/src/ui/hooks/useAuthCommand.ts
index 445b184b..283572d2 100644
--- a/packages/cli/src/ui/hooks/useAuthCommand.ts
+++ b/packages/cli/src/ui/hooks/useAuthCommand.ts
@@ -46,12 +46,16 @@ export const useAuthCommand = (
           config,
         );
       } catch (e) {
-        const errorMessage =
+        let errorMessage = `Failed to login.\nMessage: ${getErrorMessage(e)}`;
+        if (
           settings.merged.selectedAuthType ===
-          AuthType.LOGIN_WITH_GOOGLE_PERSONAL
-            ? `Failed to login. Ensure the Google account you are using is not a Workspace account and that you are not a licensed Code Assist user (see https://goo.gle/gemini-cli-auth-docs#workspace-gca).
-Message: ${getErrorMessage(e)}`
-            : `Failed to login. Message: ${getErrorMessage(e)}`;
+            AuthType.LOGIN_WITH_GOOGLE_PERSONAL &&
+          !process.env.GOOGLE_CLOUD_PROJECT
+        ) {
+          errorMessage =
+            'Failed to login. Workspace accounts and licensed Code Assist users must configure' +
+            ` GOOGLE_CLOUD_PROJECT (see https://goo.gle/gemini-cli-auth-docs#workspace-gca).\nMessage: ${getErrorMessage(e)}`;
+        }
         setAuthError(errorMessage);
         openAuthDialog();
       } finally {
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
index 3da95ffd..5ac9aaa7 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -1043,7 +1043,7 @@ describe('useGeminiStream', () => {
     it('should call parseAndFormatApiError with the correct authType on stream initialization failure', async () => {
       // 1. Setup
       const mockError = new Error('Rate limit exceeded');
-      const mockAuthType = AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE;
+      const mockAuthType = AuthType.LOGIN_WITH_GOOGLE_PERSONAL;
       mockParseAndFormatApiError.mockClear();
       mockSendMessageStream.mockReturnValue(
         (async function* () {
diff --git a/packages/cli/src/ui/utils/errorParsing.test.ts b/packages/cli/src/ui/utils/errorParsing.test.ts
index cf7234c2..72ddb0c9 100644
--- a/packages/cli/src/ui/utils/errorParsing.test.ts
+++ b/packages/cli/src/ui/utils/errorParsing.test.ts
@@ -29,12 +29,12 @@ describe('parseAndFormatApiError', () => {
     expect(result).toContain('Your request has been rate limited');
   });
 
-  it('should format a 429 API error with the enterprise message', () => {
+  it('should format a 429 API error with the personal message', () => {
     const errorMessage =
       'got status: 429 Too Many Requests. {"error":{"code":429,"message":"Rate limit exceeded","status":"RESOURCE_EXHAUSTED"}}';
     const result = parseAndFormatApiError(
       errorMessage,
-      AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE,
+      AuthType.LOGIN_WITH_GOOGLE_PERSONAL,
     );
     expect(result).toContain('[API Error: Rate limit exceeded');
     expect(result).toContain(enterpriseMessage);
diff --git a/packages/cli/src/ui/utils/errorParsing.ts b/packages/cli/src/ui/utils/errorParsing.ts
index 7935e4f3..98020ee8 100644
--- a/packages/cli/src/ui/utils/errorParsing.ts
+++ b/packages/cli/src/ui/utils/errorParsing.ts
@@ -46,7 +46,6 @@ function isStructuredError(error: unknown): error is StructuredError {
 function getRateLimitMessage(authType?: AuthType): string {
   switch (authType) {
     case AuthType.LOGIN_WITH_GOOGLE_PERSONAL:
-    case AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE:
       return RATE_LIMIT_ERROR_MESSAGE_GOOGLE;
     case AuthType.USE_GEMINI:
       return RATE_LIMIT_ERROR_MESSAGE_USE_GEMINI;
diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts
index 1ceb818b..8a3c6939 100644
--- a/packages/core/src/code_assist/codeAssist.ts
+++ b/packages/core/src/code_assist/codeAssist.ts
@@ -13,10 +13,7 @@ export async function createCodeAssistContentGenerator(
   httpOptions: HttpOptions,
   authType: AuthType,
 ): Promise<ContentGenerator> {
-  if (
-    authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE ||
-    authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL
-  ) {
+  if (authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL) {
     const authClient = await getOauthClient();
     const projectId = await setupUser(authClient);
     return new CodeAssistServer(authClient, projectId, httpOptions);
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index eb94baed..6170f319 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -505,10 +505,7 @@ export class GeminiClient {
    */
   private async handleFlashFallback(authType?: string): Promise<string | null> {
     // Only handle fallback for OAuth users
-    if (
-      authType !== AuthType.LOGIN_WITH_GOOGLE_PERSONAL &&
-      authType !== AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE
-    ) {
+    if (authType !== AuthType.LOGIN_WITH_GOOGLE_PERSONAL) {
       return null;
     }
 
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index c708dad4..7021adc2 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -36,7 +36,6 @@ export interface ContentGenerator {
 
 export enum AuthType {
   LOGIN_WITH_GOOGLE_PERSONAL = 'oauth-personal',
-  LOGIN_WITH_GOOGLE_ENTERPRISE = 'oauth-enterprise',
   USE_GEMINI = 'gemini-api-key',
   USE_VERTEX_AI = 'vertex-ai',
 }
@@ -71,14 +70,6 @@ export async function createContentGeneratorConfig(
     return contentGeneratorConfig;
   }
 
-  // if its enterprise make sure we have a cloud project
-  if (
-    authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE &&
-    !!googleCloudProject
-  ) {
-    return contentGeneratorConfig;
-  }
-
   //
   if (authType === AuthType.USE_GEMINI && geminiApiKey) {
     contentGeneratorConfig.apiKey = geminiApiKey;
@@ -118,10 +109,7 @@ export async function createContentGenerator(
       'User-Agent': `GeminiCLI/${version} (${process.platform}; ${process.arch})`,
     },
   };
-  if (
-    config.authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL ||
-    config.authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE
-  ) {
+  if (config.authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL) {
     return createCodeAssistContentGenerator(httpOptions, config.authType);
   }
 
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index 4db13852..ce5accf4 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -188,10 +188,7 @@ export class GeminiChat {
    */
   private async handleFlashFallback(authType?: string): Promise<string | null> {
     // Only handle fallback for OAuth users
-    if (
-      authType !== AuthType.LOGIN_WITH_GOOGLE_PERSONAL &&
-      authType !== AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE
-    ) {
+    if (authType !== AuthType.LOGIN_WITH_GOOGLE_PERSONAL) {
       return null;
     }
 
diff --git a/packages/core/src/utils/retry.test.ts b/packages/core/src/utils/retry.test.ts
index 031c0991..1988c02a 100644
--- a/packages/core/src/utils/retry.test.ts
+++ b/packages/core/src/utils/retry.test.ts
@@ -275,35 +275,6 @@ describe('retryWithBackoff', () => {
       expect(mockFn).toHaveBeenCalledTimes(4); // 3 initial attempts + 1 after fallback
     });
 
-    it('should trigger fallback for OAuth enterprise users after persistent 429 errors', async () => {
-      const fallbackCallback = vi.fn().mockResolvedValue('gemini-2.5-flash');
-
-      let fallbackOccurred = false;
-      const mockFn = vi.fn().mockImplementation(async () => {
-        if (!fallbackOccurred) {
-          const error: HttpError = new Error('Rate limit exceeded');
-          error.status = 429;
-          throw error;
-        }
-        return 'success';
-      });
-
-      const promise = retryWithBackoff(mockFn, {
-        maxAttempts: 3,
-        initialDelayMs: 100,
-        onPersistent429: async (authType?: string) => {
-          fallbackOccurred = true;
-          return await fallbackCallback(authType);
-        },
-        authType: 'oauth-enterprise',
-      });
-
-      await vi.runAllTimersAsync();
-
-      await expect(promise).resolves.toBe('success');
-      expect(fallbackCallback).toHaveBeenCalledWith('oauth-enterprise');
-    });
-
     it('should NOT trigger fallback for API key users', async () => {
       const fallbackCallback = vi.fn();
 
diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts
index 851db0e7..ebe18510 100644
--- a/packages/core/src/utils/retry.ts
+++ b/packages/core/src/utils/retry.ts
@@ -99,8 +99,7 @@ export async function retryWithBackoff<T>(
         if (
           consecutive429Count >= 2 &&
           onPersistent429 &&
-          (authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL ||
-            authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE)
+          authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL
         ) {
           try {
             const fallbackModel = await onPersistent429(authType);