@office-ai/aioncli-core 0.2.3 → 0.18.4

package/dist/src/core/client.js

@@ -3,97 +3,61 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import { getDirectoryContextString, getEnvironmentContext, } from '../utils/environmentContext.js';
-import { Turn, GeminiEventType, } from './turn.js';
-import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js';
-import { getResponseText, getFunctionCalls } from '../utils/generateContentResponseUtilities.js';
+import { getDirectoryContextString, getInitialChatHistory, } from '../utils/environmentContext.js';
+import { CompressionStatus } from './turn.js';
+import { Turn, GeminiEventType } from './turn.js';
+import { getCoreSystemPrompt } from './prompts.js';
 import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';
 import { reportError } from '../utils/errorReporting.js';
 import { GeminiChat } from './geminiChat.js';
 import { retryWithBackoff } from '../utils/retry.js';
 import { getErrorMessage } from '../utils/errors.js';
-import { isFunctionResponse } from '../utils/messageInspectors.js';
 import { tokenLimit } from './tokenLimits.js';
-import { AuthType, createContentGenerator, } from './contentGenerator.js';
-import { ProxyAgent, setGlobalDispatcher } from 'undici';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
+import { DEFAULT_GEMINI_FLASH_MODEL, getEffectiveModel, } from '../config/models.js';
 import { LoopDetectionService } from '../services/loopDetectionService.js';
-import { ideContext } from '../ide/ideContext.js';
-import { logChatCompression, logNextSpeakerCheck, } from '../telemetry/loggers.js';
-import { makeChatCompressionEvent, MalformedJsonResponseEvent, NextSpeakerCheckEvent, } from '../telemetry/types.js';
-import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
-function isThinkingSupported(model) {
-    if (model.startsWith('gemini-2.5'))
-        return true;
-    return false;
-}
-/**
- * Returns the index of the content after the fraction of the total characters in the history.
- *
- * Exported for testing purposes.
- */
-export function findIndexAfterFraction(history, fraction) {
-    if (fraction <= 0 || fraction >= 1) {
-        throw new Error('Fraction must be between 0 and 1');
-    }
-    const contentLengths = history.map((content) => JSON.stringify(content).length);
-    const totalCharacters = contentLengths.reduce((sum, length) => sum + length, 0);
-    const targetCharacters = totalCharacters * fraction;
-    let charactersSoFar = 0;
-    for (let i = 0; i < contentLengths.length; i++) {
-        charactersSoFar += contentLengths[i];
-        if (charactersSoFar >= targetCharacters) {
-            return i;
-        }
-    }
-    return contentLengths.length;
-}
+import { ChatCompressionService } from '../services/chatCompressionService.js';
+import { ideContextStore } from '../ide/ideContext.js';
+import { logContentRetryFailure, logNextSpeakerCheck, } from '../telemetry/loggers.js';
+import { ContentRetryFailureEvent, NextSpeakerCheckEvent, } from '../telemetry/types.js';
+import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
+import { handleFallback } from '../fallback/handler.js';
+import { debugLogger } from '../utils/debugLogger.js';
 const MAX_TURNS = 100;
-/**
- * Threshold for compression token count as a fraction of the model's token limit.
- * If the chat history exceeds this threshold, it will be compressed.
- */
-const COMPRESSION_TOKEN_THRESHOLD = 0.7;
-/**
- * The fraction of the latest chat history to keep. A value of 0.3
- * means that only the last 30% of the chat history will be kept after compression.
- */
-const COMPRESSION_PRESERVE_THRESHOLD = 0.3;
 export class GeminiClient {
     config;
     chat;
-    contentGenerator;
-    embeddingModel;
-    generateContentConfig = {
-        temperature: 0,
-        topP: 1,
-    };
     sessionTurnCount = 0;
     loopDetector;
+    compressionService;
     lastPromptId;
+    currentSequenceModel = null;
     lastSentIdeContext;
     forceFullIdeContext = true;
+    /**
+     * At any point in this conversation, was compression triggered without
+     * being forced and did it fail?
+     */
+    hasFailedCompressionAttempt = false;
     constructor(config) {
         this.config = config;
-        if (config.getProxy()) {
-            setGlobalDispatcher(new ProxyAgent(config.getProxy()));
-        }
-        this.embeddingModel = config.getEmbeddingModel();
         this.loopDetector = new LoopDetectionService(config);
+        this.compressionService = new ChatCompressionService();
         this.lastPromptId = this.config.getSessionId();
     }
-    async initialize(contentGeneratorConfig) {
-        this.contentGenerator = await createContentGenerator(contentGeneratorConfig, this.config, this.config.getSessionId());
+    updateTelemetryTokenCount() {
+        if (this.chat) {
+            uiTelemetryService.setLastPromptTokenCount(this.chat.getLastPromptTokenCount());
+        }
+    }
+    async initialize() {
         this.chat = await this.startChat();
+        this.updateTelemetryTokenCount();
     }
-    getContentGenerator() {
-        if (!this.contentGenerator) {
+    getContentGeneratorOrFail() {
+        if (!this.config.getContentGenerator()) {
             throw new Error('Content generator not initialized');
         }
-        return this.contentGenerator;
-    }
-    getUserTier() {
-        return this.contentGenerator?.userTier;
+        return this.config.getContentGenerator();
     }
     async addHistory(content) {
         this.getChat().addHistory(content);
@@ -105,32 +69,16 @@ export class GeminiClient {
         return this.chat;
     }
     isInitialized() {
-        return this.chat !== undefined && this.contentGenerator !== undefined;
+        return this.chat !== undefined;
     }
     getHistory() {
         return this.getChat().getHistory();
     }
-    setHistory(history, { stripThoughts = false } = {}) {
-        const historyToSet = stripThoughts
-            ? history.map((content) => {
-                const newContent = { ...content };
-                if (newContent.parts) {
-                    newContent.parts = newContent.parts.map((part) => {
-                        if (part &&
-                            typeof part === 'object' &&
-                            'thoughtSignature' in part) {
-                            const newPart = { ...part };
-                            delete newPart
-                                .thoughtSignature;
-                            return newPart;
-                        }
-                        return part;
-                    });
-                }
-                return newContent;
-            })
-            : history;
-        this.getChat().setHistory(historyToSet);
+    stripThoughtsFromHistory() {
+        this.getChat().stripThoughtsFromHistory();
+    }
+    setHistory(history) {
+        this.getChat().setHistory(history);
         this.forceFullIdeContext = true;
     }
     async setTools() {
@@ -141,6 +89,19 @@ export class GeminiClient {
     }
     async resetChat() {
         this.chat = await this.startChat();
+        this.updateTelemetryTokenCount();
+    }
+    async resumeChat(history, resumedSessionData) {
+        this.chat = await this.startChat(history, resumedSessionData);
+    }
+    getChatRecordingService() {
+        return this.chat?.getChatRecordingService();
+    }
+    getLoopDetectionService() {
+        return this.loopDetector;
+    }
+    getCurrentSequenceModel() {
+        return this.currentSequenceModel;
     }
     async addDirectoryContext() {
         if (!this.chat) {
@@ -151,39 +112,25 @@ export class GeminiClient {
             parts: [{ text: await getDirectoryContextString(this.config) }],
         });
     }
-    async startChat(extraHistory) {
+    async updateSystemInstruction() {
+        if (!this.isInitialized()) {
+            return;
+        }
+        const userMemory = this.config.getUserMemory();
+        const systemInstruction = getCoreSystemPrompt(this.config, userMemory);
+        this.getChat().setSystemInstruction(systemInstruction);
+    }
+    async startChat(extraHistory, resumedSessionData) {
         this.forceFullIdeContext = true;
-        const envParts = await getEnvironmentContext(this.config);
+        this.hasFailedCompressionAttempt = false;
         const toolRegistry = this.config.getToolRegistry();
         const toolDeclarations = toolRegistry.getFunctionDeclarations();
         const tools = [{ functionDeclarations: toolDeclarations }];
-        const history = [
-            {
-                role: 'user',
-                parts: envParts,
-            },
-            {
-                role: 'model',
-                parts: [{ text: 'Got it. Thanks for the context!' }],
-            },
-            ...(extraHistory ?? []),
-        ];
+        const history = await getInitialChatHistory(this.config, extraHistory);
         try {
             const userMemory = this.config.getUserMemory();
-            const systemInstruction = getCoreSystemPrompt(userMemory);
-            const generateContentConfigWithThinking = isThinkingSupported(this.config.getModel())
-                ? {
-                    ...this.generateContentConfig,
-                    thinkingConfig: {
-                        includeThoughts: true,
-                    },
-                }
-                : this.generateContentConfig;
-            return new GeminiChat(this.config, this.getContentGenerator(), {
-                systemInstruction,
-                ...generateContentConfigWithThinking,
-                tools,
-            }, history);
+            const systemInstruction = getCoreSystemPrompt(this.config, userMemory);
+            return new GeminiChat(this.config, systemInstruction, tools, history, resumedSessionData);
         }
         catch (error) {
             await reportError(error, 'Error initializing Gemini chat session.', history, 'startChat');
@@ -191,7 +138,7 @@ export class GeminiClient {
         }
     }
     getIdeContextParts(forceFullContext) {
-        const currentIdeContext = ideContext.getIdeContext();
+        const currentIdeContext = ideContextStore.get();
         if (!currentIdeContext) {
             return { contextParts: [], newIdeContext: undefined };
         }
@@ -229,7 +176,7 @@ export class GeminiClient {
                 '```',
             ];
             if (this.config.getDebugMode()) {
-                console.log(contextParts.join('\n'));
+                debugLogger.log(contextParts.join('\n'));
             }
             return {
                 contextParts,
@@ -321,7 +268,7 @@ export class GeminiClient {
                 '```',
             ];
             if (this.config.getDebugMode()) {
-                console.log(contextParts.join('\n'));
+                debugLogger.log(contextParts.join('\n'));
             }
             return {
                 contextParts,
@@ -329,10 +276,18 @@ export class GeminiClient {
             };
         }
     }
-    async *sendMessageStream(request, signal, prompt_id, turns = MAX_TURNS, originalModel) {
+    _getEffectiveModelForCurrentTurn() {
+        if (this.currentSequenceModel) {
+            return this.currentSequenceModel;
+        }
+        const configModel = this.config.getModel();
+        return getEffectiveModel(this.config.isInFallbackMode(), configModel, this.config.getPreviewFeatures());
+    }
+    async *sendMessageStream(request, signal, prompt_id, turns = MAX_TURNS, isInvalidStreamRetry = false) {
         if (this.lastPromptId !== prompt_id) {
             this.loopDetector.reset(prompt_id);
             this.lastPromptId = prompt_id;
+            this.currentSequenceModel = null;
         }
         this.sessionTurnCount++;
         if (this.config.getMaxSessionTurns() > 0 &&
@@ -345,10 +300,19 @@ export class GeminiClient {
         if (!boundedTurns) {
             return new Turn(this.getChat(), prompt_id);
         }
-        // Track the original model from the first call to detect model switching
-        const initialModel = originalModel || this.config.getModel();
-        const compressed = await this.tryCompressChat(prompt_id);
-        if (compressed) {
+        // Check for context window overflow
+        const modelForLimitCheck = this._getEffectiveModelForCurrentTurn();
+        const estimatedRequestTokenCount = Math.floor(JSON.stringify(request).length / 4);
+        const remainingTokenCount = tokenLimit(modelForLimitCheck) - this.getChat().getLastPromptTokenCount();
+        if (estimatedRequestTokenCount > remainingTokenCount * 0.95) {
+            yield {
+                type: GeminiEventType.ContextWindowWillOverflow,
+                value: { estimatedRequestTokenCount, remainingTokenCount },
+            };
+            return new Turn(this.getChat(), prompt_id);
+        }
+        const compressed = await this.tryCompressChat(prompt_id, false);
+        if (compressed.compressionStatus === CompressionStatus.COMPRESSED) {
             yield { type: GeminiEventType.ChatCompressed, value: compressed };
         }
         // Prevent context updates from being sent while a tool call is
@@ -373,205 +337,109 @@ export class GeminiClient {
             this.forceFullIdeContext = false;
         }
         const turn = new Turn(this.getChat(), prompt_id);
+        const controller = new AbortController();
+        const linkedSignal = AbortSignal.any([signal, controller.signal]);
         const loopDetected = await this.loopDetector.turnStarted(signal);
         if (loopDetected) {
             yield { type: GeminiEventType.LoopDetected };
             return turn;
         }
-        const resultStream = turn.run(request, signal);
+        const routingContext = {
+            history: this.getChat().getHistory(/*curated=*/ true),
+            request,
+            signal,
+        };
+        let modelToUse;
+        // Determine Model (Stickiness vs. Routing)
+        if (this.currentSequenceModel) {
+            modelToUse = this.currentSequenceModel;
+        }
+        else {
+            const router = await this.config.getModelRouterService();
+            const decision = await router.route(routingContext);
+            modelToUse = decision.model;
+            // Lock the model for the rest of the sequence
+            this.currentSequenceModel = modelToUse;
+            yield { type: GeminiEventType.ModelInfo, value: modelToUse };
+        }
+        const resultStream = turn.run({ model: modelToUse }, request, linkedSignal);
         for await (const event of resultStream) {
             if (this.loopDetector.addAndCheck(event)) {
                 yield { type: GeminiEventType.LoopDetected };
+                controller.abort();
                 return turn;
             }
             yield event;
+            this.updateTelemetryTokenCount();
+            if (event.type === GeminiEventType.InvalidStream) {
+                if (this.config.getContinueOnFailedApiCall()) {
+                    if (isInvalidStreamRetry) {
+                        // We already retried once, so stop here.
+                        logContentRetryFailure(this.config, new ContentRetryFailureEvent(4, // 2 initial + 2 after injections
+                        'FAILED_AFTER_PROMPT_INJECTION', modelToUse));
+                        return turn;
+                    }
+                    const nextRequest = [{ text: 'System: Please continue.' }];
+                    yield* this.sendMessageStream(nextRequest, signal, prompt_id, boundedTurns - 1, true);
+                    return turn;
+                }
+            }
+            if (event.type === GeminiEventType.Error) {
+                return turn;
+            }
         }
         if (!turn.pendingToolCalls.length && signal && !signal.aborted) {
-            // Check if model was switched during the call (likely due to quota error)
-            const currentModel = this.config.getModel();
-            if (currentModel !== initialModel) {
-                // Model was switched (likely due to quota error fallback)
-                // Don't continue with recursive call to prevent unwanted Flash execution
+            // Check if next speaker check is needed
+            if (this.config.getQuotaErrorOccurred()) {
                 return turn;
             }
             if (this.config.getSkipNextSpeakerCheck()) {
                 return turn;
             }
-            const nextSpeakerCheck = await checkNextSpeaker(this.getChat(), this, signal);
+            const nextSpeakerCheck = await checkNextSpeaker(this.getChat(), this.config.getBaseLlmClient(), signal, prompt_id);
             logNextSpeakerCheck(this.config, new NextSpeakerCheckEvent(prompt_id, turn.finishReason?.toString() || '', nextSpeakerCheck?.next_speaker || ''));
             if (nextSpeakerCheck?.next_speaker === 'model') {
                 const nextRequest = [{ text: 'Please continue.' }];
                 // This recursive call's events will be yielded out, but the final
                 // turn object will be from the top-level call.
-                yield* this.sendMessageStream(nextRequest, signal, prompt_id, boundedTurns - 1, initialModel);
+                yield* this.sendMessageStream(nextRequest, signal, prompt_id, boundedTurns - 1);
             }
         }
         return turn;
     }
-    async generateJson(contents, schema, abortSignal, model, config = {}) {
-        // Use current model from config instead of hardcoded Flash model
-        const modelToUse = model || this.config.getModel() || DEFAULT_GEMINI_FLASH_MODEL;
-        // If using OpenAI authType, use OpenAI implementation
-        console.log("my generateJson", this.config.getContentGeneratorConfig()?.authType, modelToUse);
-        if (this.config.getContentGeneratorConfig()?.authType === AuthType.USE_OPENAI) {
-            return this.openaiGenerateJson(contents, schema, abortSignal, modelToUse, config);
-        }
-        try {
-            const userMemory = this.config.getUserMemory();
-            const systemInstruction = getCoreSystemPrompt(userMemory);
-            const requestConfig = {
-                abortSignal,
-                ...this.generateContentConfig,
-                ...config,
-            };
-            const apiCall = () => this.getContentGenerator().generateContent({
-                model: modelToUse,
-                config: {
-                    ...requestConfig,
-                    systemInstruction,
-                    responseJsonSchema: schema,
-                    responseMimeType: 'application/json',
-                },
-                contents,
-            }, this.lastPromptId);
-            const result = await retryWithBackoff(apiCall, {
-                onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
-                authType: this.config.getContentGeneratorConfig()?.authType,
-            });
-            let text = getResponseText(result);
-            if (!text) {
-                const error = new Error('API returned an empty response for generateJson.');
-                await reportError(error, 'Error in generateJson: API returned an empty response.', contents, 'generateJson-empty-response');
-                throw error;
-            }
-            const prefix = '```json';
-            const suffix = '```';
-            if (text.startsWith(prefix) && text.endsWith(suffix)) {
-                ClearcutLogger.getInstance(this.config)?.logMalformedJsonResponseEvent(new MalformedJsonResponseEvent(modelToUse));
-                text = text
-                    .substring(prefix.length, text.length - suffix.length)
-                    .trim();
-            }
-            try {
-                return JSON.parse(text);
-            }
-            catch (parseError) {
-                await reportError(parseError, 'Failed to parse JSON response from generateJson.', {
-                    responseTextFailedToParse: text,
-                    originalRequestContents: contents,
-                }, 'generateJson-parse');
-                throw new Error(`Failed to parse API response as JSON: ${getErrorMessage(parseError)}`);
-            }
-        }
-        catch (error) {
-            if (abortSignal.aborted) {
-                throw error;
-            }
-            // Avoid double reporting for the empty response case handled above
-            if (error instanceof Error &&
-                error.message === 'API returned an empty response for generateJson.') {
-                throw error;
-            }
-            await reportError(error, 'Error generating JSON content via API.', contents, 'generateJson-api');
-            throw new Error(`Failed to generate JSON content: ${getErrorMessage(error)}`);
-        }
-    }
-    async openaiGenerateJson(contents, schema, abortSignal, model, config = {}) {
+    async generateContent(modelConfigKey, contents, abortSignal) {
+        const desiredModelConfig = this.config.modelConfigService.getResolvedConfig(modelConfigKey);
+        let { model: currentAttemptModel, generateContentConfig: currentAttemptGenerateContentConfig, } = desiredModelConfig;
+        const fallbackModelConfig = this.config.modelConfigService.getResolvedConfig({
+            ...modelConfigKey,
+            model: DEFAULT_GEMINI_FLASH_MODEL,
+        });
         try {
             const userMemory = this.config.getUserMemory();
-            const systemInstruction = getCoreSystemPrompt(userMemory);
-            const requestConfig = {
-                abortSignal,
-                ...this.generateContentConfig,
-                ...config,
-            };
-            // Create function declaration for OpenAI-style response
-            const functionDeclaration = {
-                name: 'respond_in_schema',
-                description: 'Provide the response in provided schema',
-                parameters: schema,
-            };
-            const tools = [
-                {
-                    functionDeclarations: [functionDeclaration],
-                },
-            ];
-            const apiCall = () => this.getContentGenerator().generateContent({
-                model,
-                config: {
-                    ...requestConfig,
+            const systemInstruction = getCoreSystemPrompt(this.config, userMemory);
+            const apiCall = () => {
+                const modelConfigToUse = this.config.isInFallbackMode()
+                    ? fallbackModelConfig
+                    : desiredModelConfig;
+                currentAttemptModel = modelConfigToUse.model;
+                currentAttemptGenerateContentConfig =
+                    modelConfigToUse.generateContentConfig;
+                const requestConfig = {
+                    ...currentAttemptGenerateContentConfig,
+                    abortSignal,
                     systemInstruction,
-                    tools,
-                },
-                contents,
-            }, this.lastPromptId);
-            const result = await retryWithBackoff(apiCall, {
-                onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
-                authType: this.config.getContentGeneratorConfig()?.authType,
-            });
-            // Try to extract function call arguments
-            const functionCalls = getFunctionCalls(result);
-            if (functionCalls && functionCalls.length > 0) {
-                const functionCall = functionCalls.find((call) => call.name === 'respond_in_schema');
-                if (functionCall && functionCall.args) {
-                    return functionCall.args;
-                }
-            }
-            // Fallback: try to parse response text as JSON
-            let text = getResponseText(result);
-            if (!text) {
-                const error = new Error('API returned an empty response for openaiGenerateJson.');
-                await reportError(error, 'Error in openaiGenerateJson: API returned an empty response.', contents, 'openaiGenerateJson-empty-response');
-                throw error;
-            }
-            // Handle markdown wrapped JSON
-            const prefix = '```json';
-            const suffix = '```';
-            if (text.startsWith(prefix) && text.endsWith(suffix)) {
-                text = text
-                    .substring(prefix.length, text.length - suffix.length)
-                    .trim();
-            }
-            try {
-                return JSON.parse(text);
-            }
-            catch (parseError) {
-                await reportError(parseError, 'Failed to parse JSON response from openaiGenerateJson.', {
-                    responseTextFailedToParse: text,
-                    originalRequestContents: contents,
-                }, 'openaiGenerateJson-parse');
-                throw new Error(`Failed to parse API response as JSON: ${getErrorMessage(parseError)}`);
-            }
-        }
-        catch (error) {
-            if (abortSignal.aborted) {
-                throw error;
-            }
-            await reportError(error, 'Error generating JSON content via OpenAI API.', contents, 'openaiGenerateJson-api');
-            throw new Error(`Failed to generate JSON content: ${getErrorMessage(error)}`);
-        }
-    }
-    async generateContent(contents, generationConfig, abortSignal, model) {
-        const modelToUse = model ?? this.config.getModel();
-        const configToUse = {
-            ...this.generateContentConfig,
-            ...generationConfig,
-        };
-        try {
-            const userMemory = this.config.getUserMemory();
-            const systemInstruction = getCoreSystemPrompt(userMemory);
-            const requestConfig = {
-                abortSignal,
-                ...configToUse,
-                systemInstruction,
+                };
+                return this.getContentGeneratorOrFail().generateContent({
+                    model: currentAttemptModel,
+                    config: requestConfig,
+                    contents,
+                }, this.lastPromptId);
             };
-            const apiCall = () => this.getContentGenerator().generateContent({
-                model: modelToUse,
-                config: requestConfig,
-                contents,
-            }, this.lastPromptId);
+            const onPersistent429Callback = async (authType, error) => 
+            // Pass the captured model to the centralized handler.
+            await handleFallback(this.config, currentAttemptModel, authType, error);
             const result = await retryWithBackoff(apiCall, {
-                onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
+                onPersistent429: onPersistent429Callback,
                 authType: this.config.getContentGeneratorConfig()?.authType,
             });
             return result;
@@ -580,198 +448,31 @@ export class GeminiClient {
             if (abortSignal.aborted) {
                 throw error;
             }
-            await reportError(error, `Error generating content via API with model ${modelToUse}.`, {
+            await reportError(error, `Error generating content via API with model ${currentAttemptModel}.`, {
                 requestContents: contents,
-                requestConfig: configToUse,
+                requestConfig: currentAttemptGenerateContentConfig,
             }, 'generateContent-api');
-            throw new Error(`Failed to generate content with model ${modelToUse}: ${getErrorMessage(error)}`);
-        }
-    }
-    async generateEmbedding(texts) {
-        if (!texts || texts.length === 0) {
-            return [];
+            throw new Error(`Failed to generate content with model ${currentAttemptModel}: ${getErrorMessage(error)}`);
         }
-        const embedModelParams = {
-            model: this.embeddingModel,
-            contents: texts,
-        };
-        const embedContentResponse = await this.getContentGenerator().embedContent(embedModelParams);
-        if (!embedContentResponse.embeddings ||
-            embedContentResponse.embeddings.length === 0) {
-            throw new Error('No embeddings found in API response.');
-        }
-        if (embedContentResponse.embeddings.length !== texts.length) {
-            throw new Error(`API returned a mismatched number of embeddings. Expected ${texts.length}, got ${embedContentResponse.embeddings.length}.`);
-        }
-        return embedContentResponse.embeddings.map((embedding, index) => {
-            const values = embedding.values;
-            if (!values || values.length === 0) {
-                throw new Error(`API returned an empty embedding for input text at index ${index}: "${texts[index]}"`);
-            }
-            return values;
-        });
     }
     async tryCompressChat(prompt_id, force = false) {
-        const curatedHistory = this.getChat().getHistory(true);
-        // Regardless of `force`, don't do anything if the history is empty.
-        if (curatedHistory.length === 0) {
-            return null;
-        }
-        const model = this.config.getModel();
-        const { totalTokens: originalTokenCount } = await this.getContentGenerator().countTokens({
-            model,
-            contents: curatedHistory,
-        });
-        if (originalTokenCount === undefined) {
-            console.warn(`Could not determine token count for model ${model}.`);
-            return null;
-        }
-        const contextPercentageThreshold = this.config.getChatCompression()?.contextPercentageThreshold;
-        // Don't compress if not forced and we are under the limit.
-        if (!force) {
-            const threshold = contextPercentageThreshold ?? COMPRESSION_TOKEN_THRESHOLD;
-            if (originalTokenCount < threshold * tokenLimit(model)) {
-                return null;
-            }
-        }
-        let compressBeforeIndex = findIndexAfterFraction(curatedHistory, 1 - COMPRESSION_PRESERVE_THRESHOLD);
-        // Find the first user message after the index. This is the start of the next turn.
-        while (compressBeforeIndex < curatedHistory.length &&
-            (curatedHistory[compressBeforeIndex]?.role === 'model' ||
-                isFunctionResponse(curatedHistory[compressBeforeIndex]))) {
-            compressBeforeIndex++;
-        }
-        const historyToCompress = curatedHistory.slice(0, compressBeforeIndex);
-        const historyToKeep = curatedHistory.slice(compressBeforeIndex);
-        this.getChat().setHistory(historyToCompress);
-        const { text: summary } = await this.getChat().sendMessage({
-            message: {
-                text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.',
-            },
-            config: {
-                systemInstruction: { text: getCompressionPrompt() },
-            },
-        }, prompt_id);
-        this.chat = await this.startChat([
-            {
-                role: 'user',
-                parts: [{ text: summary }],
-            },
-            {
-                role: 'model',
-                parts: [{ text: 'Got it. Thanks for the additional context!' }],
-            },
-            ...historyToKeep,
-        ]);
-        this.forceFullIdeContext = true;
-        const { totalTokens: newTokenCount } = await this.getContentGenerator().countTokens({
-            // model might change after calling `sendMessage`, so we get the newest value from config
-            model: this.config.getModel(),
-            contents: this.getChat().getHistory(),
-        });
-        if (newTokenCount === undefined) {
-            console.warn('Could not determine compressed history token count.');
-            return null;
-        }
-        logChatCompression(this.config, makeChatCompressionEvent({
-            tokens_before: originalTokenCount,
-            tokens_after: newTokenCount,
-        }));
-        return {
-            originalTokenCount,
-            newTokenCount,
-        };
-    }
-    /**
-     * [PATCH:API_KEY_ROTATION] Tries to rotate API key for GEMINI/OPENAI API key authentication modes.
-     * Returns the current model name if key was rotated, null otherwise.
-     */
-    async tryRotateApiKey(authType) {
-        // Support both GEMINI and OPENAI API key modes
-        let envKey;
-        if (authType === AuthType.USE_GEMINI) {
-            envKey = 'GEMINI_API_KEY';
-        }
-        else if (authType === AuthType.USE_OPENAI) {
-            envKey = 'OPENAI_API_KEY';
-        }
-        else {
-            return null; // Not a supported API key mode
-        }
-        const newApiKey = process.env[envKey]?.trim();
-        if (!newApiKey) {
-            return null;
-        }
-        const currentConfig = this.config.getContentGeneratorConfig();
-        if (!currentConfig) {
-            return null;
-        }
-        if (currentConfig.apiKey === newApiKey) {
-            return null; // Same key, no rotation needed
-        }
-        try {
-            // Update the API key in config first
-            currentConfig.apiKey = newApiKey;
-            // Create new content generator with the new API key
-            const newContentGenerator = await createContentGenerator(currentConfig, this.config, this.config.getSessionId());
-            // Update our content generator
-            this.contentGenerator = newContentGenerator;
-            const modelName = this.config.getModel();
-            return modelName;
-        }
-        catch (_error) {
-            return null;
-        }
-    }
-    // [/PATCH:API_KEY_ROTATION]
-    /**
-     * Handles falling back to Flash model when persistent 429 errors occur for OAuth users.
-     * Uses a fallback handler if provided by the config; otherwise, returns null.
-     */
-    async handleFlashFallback(authType, error) {
-        // First try to rotate API key for GEMINI/OPENAI API key modes
-        const rotateResult = await this.tryRotateApiKey(authType);
-        if (rotateResult) {
-            return rotateResult;
-        }
-        const currentModel = this.config.getModel();
-        const fallbackModel = DEFAULT_GEMINI_FLASH_MODEL;
-        // Don't fallback if already using Flash model
-        if (currentModel === fallbackModel) {
-            return null;
-        }
-        // Check if config has a fallback handler (set by CLI package)
-        const fallbackHandler = this.config.flashFallbackHandler;
-        if (typeof fallbackHandler === 'function') {
-            try {
-                const accepted = await fallbackHandler(currentModel, fallbackModel, error);
-                if (accepted !== false && accepted !== null) {
-                    // Only OAuth users need model fallback, API key users rely on key rotation
-                    if (authType === AuthType.LOGIN_WITH_GOOGLE) {
-                        this.config.setModel(fallbackModel);
-                        this.config.setFallbackMode(true);
-                        return fallbackModel;
-                    }
-                    else {
-                        const currentModel = this.config.getModel();
-                        return currentModel; // Continue with original model, rely on API key rotation
-                    }
-                }
-                // Check if the model was switched manually in the handler
-                const newModel = this.config.getModel();
-                if (newModel === fallbackModel) {
-                    return null; // Model was switched but don't continue with current prompt
-                }
-            }
-            catch (handlerError) {
-                console.warn('Flash fallback handler failed:', handlerError);
-            }
-        }
-        return null;
+        // If the model is 'auto', we will use a placeholder model to check.
+        // Compression occurs before we choose a model, so calling `count_tokens`
+        // before the model is chosen would result in an error.
+        const model = this._getEffectiveModelForCurrentTurn();
+        const { newHistory, info } = await this.compressionService.compress(this.getChat(), prompt_id, force, model, this.config, this.hasFailedCompressionAttempt);
+        if (info.compressionStatus ===
+            CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT) {
+            this.hasFailedCompressionAttempt = !force && true;
+        }
+        else if (info.compressionStatus === CompressionStatus.COMPRESSED) {
+            if (newHistory) {
+                this.chat = await this.startChat(newHistory);
+                this.updateTelemetryTokenCount();
+                this.forceFullIdeContext = true;
+            }
+        }
+        return info;
     }
 }
-export const TEST_ONLY = {
-    COMPRESSION_PRESERVE_THRESHOLD,
-    COMPRESSION_TOKEN_THRESHOLD,
-};
 //# sourceMappingURL=client.js.map