Simplify streaming code for code assist server (#4619)

1 files changed, 10 insertions, 2 deletions

diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts
index 3c7b81b0..8831d24b 100644
--- a/packages/core/src/code_assist/setup.ts
+++ b/packages/core/src/code_assist/setup.ts
@@ -22,12 +22,17 @@ export class ProjectIdRequiredError extends Error {
   }
 }
 
+export interface UserData {
+  projectId: string;
+  userTier: UserTierId;
+}
+
 /**
  *
  * @param projectId the user's project id, if any
  * @returns the user's actual project id
  */
-export async function setupUser(client: OAuth2Client): Promise<string> {
+export async function setupUser(client: OAuth2Client): Promise<UserData> {
   let projectId = process.env.GOOGLE_CLOUD_PROJECT || undefined;
   const caServer = new CodeAssistServer(client, projectId);
 
@@ -64,7 +69,10 @@ export async function setupUser(client: OAuth2Client): Promise<string> {
     await new Promise((f) => setTimeout(f, 5000));
     lroRes = await caServer.onboardUser(onboardReq);
   }
-  return lroRes.response?.cloudaicompanionProject?.id || '';
+  return {
+    projectId: lroRes.response?.cloudaicompanionProject?.id || '',
+    userTier: tier.id,
+  };
 }
 
 function getOnboardTier(res: LoadCodeAssistResponse): GeminiUserTier {