Simplify streaming code for code assist server (#4619)

8 files changed, 30 insertions, 115 deletions

diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts
index 23dfe403..f9c6a7a3 100644
--- a/packages/core/src/code_assist/codeAssist.ts
+++ b/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}`);
diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts
index fe8661f1..a6cf6af7 100644
--- a/packages/core/src/code_assist/server.ts
+++ b/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}`;
diff --git a/packages/core/src/code_assist/setup.test.ts b/packages/core/src/code_assist/setup.test.ts
index 479abae0..6db5fd88 100644
--- a/packages/core/src/code_assist/setup.test.ts
+++ b/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 () => {
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 {
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 9528f648..6a3a18b6 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -11,7 +11,6 @@ import {
   ContentGeneratorConfig,
   createContentGeneratorConfig,
 } from '../core/contentGenerator.js';
-import { UserTierId } from '../code_assist/types.js';
 import { ToolRegistry } from '../tools/tool-registry.js';
 import { LSTool } from '../tools/ls.js';
 import { ReadFileTool } from '../tools/read-file.js';
@@ -359,14 +358,6 @@ export class Config {
     return this.quotaErrorOccurred;
   }
 
-  async getUserTier(): Promise<UserTierId | undefined> {
-    if (!this.geminiClient) {
-      return undefined;
-    }
-    const generator = this.geminiClient.getContentGenerator();
-    return await generator.getTier?.();
-  }
-
   getEmbeddingModel(): string {
     return this.embeddingModel;
   }
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index 969e1da6..61195e2f 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -22,6 +22,7 @@ import {
   ChatCompressionInfo,
 } from './turn.js';
 import { Config } from '../config/config.js';
+import { UserTierId } from '../code_assist/types.js';
 import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js';
 import { ReadManyFilesTool } from '../tools/read-many-files.js';
 import { getResponseText } from '../utils/generateContentResponseUtilities.js';
@@ -130,6 +131,10 @@ export class GeminiClient {
     return this.contentGenerator;
   }
 
+  getUserTier(): UserTierId | undefined {
+    return this.contentGenerator?.userTier;
+  }
+
   async addHistory(content: Content) {
     this.getChat().addHistory(content);
   }
diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index b381de8e..721a16d7 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -35,7 +35,7 @@ export interface ContentGenerator {
 
   embedContent(request: EmbedContentParameters): Promise<EmbedContentResponse>;
 
-  getTier?(): Promise<UserTierId | undefined>;
+  userTier?: UserTierId;
 }
 
 export enum AuthType {
diff --git a/packages/core/src/utils/quotaErrorDetection.ts b/packages/core/src/utils/quotaErrorDetection.ts
index b07309cd..6fe9b312 100644
--- a/packages/core/src/utils/quotaErrorDetection.ts
+++ b/packages/core/src/utils/quotaErrorDetection.ts
@@ -68,10 +68,6 @@ export function isProQuotaExceededError(error: unknown): boolean {
       };
     };
     if (gaxiosError.response && gaxiosError.response.data) {
-      console.log(
-        '[DEBUG] isProQuotaExceededError - checking response data:',
-        gaxiosError.response.data,
-      );
       if (typeof gaxiosError.response.data === 'string') {
         return checkMessage(gaxiosError.response.data);
       }
@@ -87,11 +83,6 @@ export function isProQuotaExceededError(error: unknown): boolean {
       }
     }
   }
-
-  console.log(
-    '[DEBUG] isProQuotaExceededError - no matching error format for:',
-    error,
-  );
   return false;
 }