diff options
Diffstat (limited to 'packages/core/src/code_assist')
| -rw-r--r-- | packages/core/src/code_assist/codeAssist.ts | 10 | ||||
| -rw-r--r-- | packages/core/src/code_assist/server.ts | 93 | ||||
| -rw-r--r-- | packages/core/src/code_assist/setup.test.ts | 5 | ||||
| -rw-r--r-- | packages/core/src/code_assist/setup.ts | 12 |
4 files changed, 24 insertions, 96 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 { |
