1 files changed, 26 insertions, 32 deletions

diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts
index d8a6dd45..a3162c81 100644
--- a/packages/core/src/code_assist/setup.ts
+++ b/packages/core/src/code_assist/setup.ts
@@ -4,52 +4,46 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { ClientMetadata, OnboardUserRequest } from './types.js';
+import { CcpaServer } from './ccpaServer.js';
 import { OAuth2Client } from 'google-auth-library';
 
-import { ClientMetadata } from './metadata.js';
-import { doLoadCodeAssist, LoadCodeAssistResponse } from './load.js';
-import { doGCALogin } from './login.js';
-import {
-  doOnboardUser,
-  LongrunningOperationResponse,
-  OnboardUserRequest,
-} from './onboard.js';
-
-export async function doSetup(): Promise<string> {
-  const oauth2Client: OAuth2Client = await doGCALogin();
+/**
+ *
+ * @param projectId the user's project id, if any
+ * @returns the user's actual project id
+ */
+export async function setupUser(
+  oAuth2Client: OAuth2Client,
+  projectId?: string,
+): Promise<string> {
+  const ccpaServer: CcpaServer = new CcpaServer(oAuth2Client, projectId);
   const clientMetadata: ClientMetadata = {
     ideType: 'IDE_UNSPECIFIED',
-    ideVersion: null,
-    pluginVersion: null,
     platform: 'PLATFORM_UNSPECIFIED',
-    updateChannel: null,
-    duetProject: 'aipp-internal-testing',
     pluginType: 'GEMINI',
-    ideName: null,
   };
+  if (process.env.GOOGLE_CLOUD_PROJECT) {
+    clientMetadata.duetProject = process.env.GOOGLE_CLOUD_PROJECT;
+  }
 
-  // Call LoadCodeAssist.
-  const loadCodeAssistRes: LoadCodeAssistResponse = await doLoadCodeAssist(
-    {
-      cloudaicompanionProject: 'aipp-internal-testing',
-      metadata: clientMetadata,
-    },
-    oauth2Client,
-  );
+  // TODO: Support Free Tier user without projectId.
+  const loadRes = await ccpaServer.loadCodeAssist({
+    cloudaicompanionProject: process.env.GOOGLE_CLOUD_PROJECT,
+    metadata: clientMetadata,
+  });
 
-  // Call OnboardUser until long running operation is complete.
-  const onboardUserReq: OnboardUserRequest = {
+  const onboardReq: OnboardUserRequest = {
     tierId: 'legacy-tier',
-    cloudaicompanionProject: loadCodeAssistRes.cloudaicompanionProject || '',
+    cloudaicompanionProject: loadRes.cloudaicompanionProject || '',
     metadata: clientMetadata,
   };
-  let lroRes: LongrunningOperationResponse = await doOnboardUser(
-    onboardUserReq,
-    oauth2Client,
-  );
+
+  // Poll onboardUser until long running operation is complete.
+  let lroRes = await ccpaServer.onboardUser(onboardReq);
   while (!lroRes.done) {
     await new Promise((f) => setTimeout(f, 5000));
-    lroRes = await doOnboardUser(onboardUserReq, oauth2Client);
+    lroRes = await ccpaServer.onboardUser(onboardReq);
   }
 
   return lroRes.response?.cloudaicompanionProject?.id || '';