Simplify streaming code for code assist server (#4619)

packages/core/src/code_assist/codeAssist.ts

@@ -21,8 +21,14 @@ export async function createCodeAssistContentGenerator(
     authType === AuthType.CLOUD_SHELL
   ) {
     const authClient = await getOauthClient(authType, config);
-    const projectId = await setupUser(authClient);
-    return new CodeAssistServer(authClient, projectId, httpOptions, sessionId);
+    const userData = await setupUser(authClient);
+    return new CodeAssistServer(
+      authClient,
+      userData.projectId,
+      httpOptions,
+      sessionId,
+      userData.userTier,
+    );
   }
 
   throw new Error(`Unsupported authType: ${authType}`);

packages/core/src/code_assist/server.ts

@@ -32,23 +32,6 @@ import {
   toCountTokenRequest,
   toGenerateContentRequest,
 } from './converter.js';
-import { Readable } from 'node:stream';
-
-interface ErrorData {
-  error?: {
-    message?: string;
-  };
-}
-
-interface GaxiosResponse {
-  status: number;
-  data: unknown;
-}
-
-interface StreamError extends Error {
-  status?: number;
-  response?: GaxiosResponse;
-}
 
 /** HTTP options to be used in each of the requests. */
 export interface HttpOptions {
@@ -60,13 +43,12 @@ export const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com';
 export const CODE_ASSIST_API_VERSION = 'v1internal';
 
 export class CodeAssistServer implements ContentGenerator {
-  private userTier: UserTierId | undefined = undefined;
-
   constructor(
     readonly client: OAuth2Client,
     readonly projectId?: string,
     readonly httpOptions: HttpOptions = {},
     readonly sessionId?: string,
+    readonly userTier?: UserTierId,
   ) {}
 
   async generateContentStream(
@@ -196,45 +178,8 @@ export class CodeAssistServer implements ContentGenerator {
     });
 
     return (async function* (): AsyncGenerator<T> {
-      // Convert ReadableStream to Node.js stream if needed
-      let nodeStream: NodeJS.ReadableStream;
-
-      if (res.data instanceof ReadableStream) {
-        // Convert Web ReadableStream to Node.js Readable stream
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        nodeStream = Readable.fromWeb(res.data as any);
-      } else if (
-        res.data &&
-        typeof (res.data as NodeJS.ReadableStream).on === 'function'
-      ) {
-        // Already a Node.js stream
-        nodeStream = res.data as NodeJS.ReadableStream;
-      } else {
-        // If res.data is not a stream, it might be an error response
-        // Try to extract error information from the response
-        let errorMessage =
-          'Response data is not a readable stream. This may indicate a server error or quota issue.';
-
-        if (res.data && typeof res.data === 'object') {
-          // Check if this is an error response with error details
-          const errorData = res.data as ErrorData;
-          if (errorData.error?.message) {
-            errorMessage = errorData.error.message;
-          } else if (typeof errorData === 'string') {
-            errorMessage = errorData;
-          }
-        }
-
-        // Create an error that looks like a quota error if it contains quota information
-        const error: StreamError = new Error(errorMessage);
-        // Add status and response properties so it can be properly handled by retry logic
-        error.status = res.status;
-        error.response = res;
-        throw error;
-      }
-
       const rl = readline.createInterface({
-        input: nodeStream,
+        input: res.data as NodeJS.ReadableStream,
         crlfDelay: Infinity, // Recognizes '\r\n' and '\n' as line breaks
       });
 
@@ -256,40 +201,6 @@ export class CodeAssistServer implements ContentGenerator {
     })();
   }
 
-  async getTier(): Promise<UserTierId | undefined> {
-    if (this.userTier === undefined) {
-      await this.detectUserTier();
-    }
-    return this.userTier;
-  }
-
-  private async detectUserTier(): Promise<void> {
-    try {
-      // Reset user tier when detection runs
-      this.userTier = undefined;
-
-      // Only attempt tier detection if we have a project ID
-      if (this.projectId) {
-        const loadRes = await this.loadCodeAssist({
-          cloudaicompanionProject: this.projectId,
-          metadata: {
-            ideType: 'IDE_UNSPECIFIED',
-            platform: 'PLATFORM_UNSPECIFIED',
-            pluginType: 'GEMINI',
-            duetProject: this.projectId,
-          },
-        });
-        if (loadRes.currentTier) {
-          this.userTier = loadRes.currentTier.id;
-        }
-      }
-    } catch (error) {
-      // Silently fail - this is not critical functionality
-      // We'll default to FREE tier behavior if tier detection fails
-      console.debug('User tier detection failed:', error);
-    }
-  }
-
   getMethodUrl(method: string): string {
     const endpoint = process.env.CODE_ASSIST_ENDPOINT ?? CODE_ASSIST_ENDPOINT;
     return `${endpoint}/${CODE_ASSIST_API_VERSION}:${method}`;

packages/core/src/code_assist/setup.test.ts

@@ -65,7 +65,10 @@ describe('setupUser', () => {
       expect.any(Object),
       undefined,
     );
-    expect(projectId).toBe('server-project');
+    expect(projectId).toEqual({
+      projectId: 'server-project',
+      userTier: 'standard-tier',
+    });
   });
 
   it('should throw ProjectIdRequiredError when no project ID is available', async () => {

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 {