From 04518b52c0ddcd5ae1192763c55e472add218b3c Mon Sep 17 00:00:00 2001
From: matt korwel <matt.korwel@gmail.com>
Date: Thu, 19 Jun 2025 16:52:22 -0700
Subject: Auth First Run (#1207)

Co-authored-by: Tommaso Sciortino <sciortino@gmail.com>
Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
---
 packages/core/src/code_assist/codeAssist.ts | 16 ++++++++++++----
 packages/core/src/code_assist/errors.ts     | 13 +++++++++++++
 packages/core/src/code_assist/oauth2.ts     |  8 ++++++++
 packages/core/src/code_assist/server.ts     |  4 ++--
 4 files changed, 35 insertions(+), 6 deletions(-)
 create mode 100644 packages/core/src/code_assist/errors.ts

(limited to 'packages/core/src/code_assist')

diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts
index 4b46f8b9..1ceb818b 100644
--- a/packages/core/src/code_assist/codeAssist.ts
+++ b/packages/core/src/code_assist/codeAssist.ts
@@ -4,15 +4,23 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ContentGenerator } from '../core/contentGenerator.js';
+import { AuthType, ContentGenerator } from '../core/contentGenerator.js';
 import { getOauthClient } from './oauth2.js';
 import { setupUser } from './setup.js';
 import { CodeAssistServer, HttpOptions } from './server.js';
 
 export async function createCodeAssistContentGenerator(
   httpOptions: HttpOptions,
+  authType: AuthType,
 ): Promise<ContentGenerator> {
-  const authClient = await getOauthClient();
-  const projectId = await setupUser(authClient);
-  return new CodeAssistServer(authClient, projectId, httpOptions);
+  if (
+    authType === AuthType.LOGIN_WITH_GOOGLE_ENTERPRISE ||
+    authType === AuthType.LOGIN_WITH_GOOGLE_PERSONAL
+  ) {
+    const authClient = await getOauthClient();
+    const projectId = await setupUser(authClient);
+    return new CodeAssistServer(authClient, projectId, httpOptions);
+  }
+
+  throw new Error(`Unsupported authType: ${authType}`);
 }
diff --git a/packages/core/src/code_assist/errors.ts b/packages/core/src/code_assist/errors.ts
new file mode 100644
index 00000000..f870ab5c
--- /dev/null
+++ b/packages/core/src/code_assist/errors.ts
@@ -0,0 +1,13 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { GaxiosError } from 'gaxios';
+
+export function isAuthError(error: unknown): boolean {
+  return (
+    error instanceof GaxiosError && error.response?.data?.error?.code === 401
+  );
+}
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index 3fbbd896..68f6e137 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -192,3 +192,11 @@ async function cacheCredentials(credentials: Credentials) {
 function getCachedCredentialPath(): string {
   return path.join(os.homedir(), GEMINI_DIR, CREDENTIAL_FILENAME);
 }
+
+export async function clearCachedCredentialFile() {
+  try {
+    await fs.rm(getCachedCredentialPath());
+  } catch (_) {
+    /* empty */
+  }
+}
diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts
index 6663eda3..4f8bb643 100644
--- a/packages/core/src/code_assist/server.ts
+++ b/packages/core/src/code_assist/server.ts
@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { OAuth2Client } from 'google-auth-library';
+import { AuthClient } from 'google-auth-library';
 import {
   LoadCodeAssistResponse,
   LoadCodeAssistRequest,
@@ -45,7 +45,7 @@ export const CODE_ASSIST_API_VERSION = 'v1internal';
 
 export class CodeAssistServer implements ContentGenerator {
   constructor(
-    readonly auth: OAuth2Client,
+    readonly auth: AuthClient,
     readonly projectId?: string,
     readonly httpOptions: HttpOptions = {},
   ) {}
-- 
cgit v1.2.3