summaryrefslogtreecommitdiff
path: root/packages/core
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core')
-rw-r--r--packages/core/src/code_assist/server.ts37
-rw-r--r--packages/core/src/config/config.ts9
-rw-r--r--packages/core/src/core/contentGenerator.ts3
3 files changed, 49 insertions, 0 deletions
diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts
index 01fd2462..fe8661f1 100644
--- a/packages/core/src/code_assist/server.ts
+++ b/packages/core/src/code_assist/server.ts
@@ -23,6 +23,7 @@ import {
} from '@google/genai';
import * as readline from 'readline';
import { ContentGenerator } from '../core/contentGenerator.js';
+import { UserTierId } from './types.js';
import {
CaCountTokenResponse,
CaGenerateContentResponse,
@@ -59,6 +60,8 @@ 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,
@@ -253,6 +256,40 @@ 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/config/config.ts b/packages/core/src/config/config.ts
index 12767133..dc85c61a 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -11,6 +11,7 @@ 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';
@@ -323,6 +324,14 @@ 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/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts
index fee10fad..d3434c23 100644
--- a/packages/core/src/core/contentGenerator.ts
+++ b/packages/core/src/core/contentGenerator.ts
@@ -17,6 +17,7 @@ import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
import { DEFAULT_GEMINI_MODEL } from '../config/models.js';
import { Config } from '../config/config.js';
import { getEffectiveModel } from './modelCheck.js';
+import { UserTierId } from '../code_assist/types.js';
/**
* Interface abstracting the core functionalities for generating content and counting tokens.
@@ -33,6 +34,8 @@ export interface ContentGenerator {
countTokens(request: CountTokensParameters): Promise<CountTokensResponse>;
embedContent(request: EmbedContentParameters): Promise<EmbedContentResponse>;
+
+ getTier?(): Promise<UserTierId | undefined>;
}
export enum AuthType {