@office-ai/aioncli-core 0.2.2 → 0.8.1

package/dist/src/core/client.js

@@ -4,49 +4,67 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { getDirectoryContextString, getEnvironmentContext, } from '../utils/environmentContext.js';
-import { Turn, GeminiEventType, } from './turn.js';
+import { CompressionStatus } from './turn.js';
+import { Turn, GeminiEventType } from './turn.js';
 import { getCoreSystemPrompt, getCompressionPrompt } from './prompts.js';
-import { getResponseText, getFunctionCalls } from '../utils/generateContentResponseUtilities.js';
+import { getResponseText } from '../utils/partUtils.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, DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_MODEL_AUTO, DEFAULT_THINKING_MODE, getEffectiveModel, } from '../config/models.js';
 import { LoopDetectionService } from '../services/loopDetectionService.js';
-import { ideContext } from '../ide/ideContext.js';
+import { ideContextStore } 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;
+import { makeChatCompressionEvent, NextSpeakerCheckEvent, } from '../telemetry/types.js';
+import { handleFallback } from '../fallback/handler.js';
+import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
+export function isThinkingSupported(model) {
+    return model.startsWith('gemini-2.5') || model === DEFAULT_GEMINI_MODEL_AUTO;
+}
+export function isThinkingDefault(model) {
+    if (model.startsWith('gemini-2.5-flash-lite')) {
+        return false;
+    }
+    return model.startsWith('gemini-2.5') || model === DEFAULT_GEMINI_MODEL_AUTO;
 }
 /**
- * Returns the index of the content after the fraction of the total characters in the history.
+ * Returns the index of the oldest item to keep when compressing. May return
+ * contents.length which indicates that everything should be compressed.
  *
  * Exported for testing purposes.
  */
-export function findIndexAfterFraction(history, fraction) {
+export function findCompressSplitPoint(contents, 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;
-        }
+    const charCounts = contents.map((content) => JSON.stringify(content).length);
+    const totalCharCount = charCounts.reduce((a, b) => a + b, 0);
+    const targetCharCount = totalCharCount * fraction;
+    let lastSplitPoint = 0; // 0 is always valid (compress nothing)
+    let cumulativeCharCount = 0;
+    for (let i = 0; i < contents.length; i++) {
+        const content = contents[i];
+        if (content.role === 'user' &&
+            !content.parts?.some((part) => !!part.functionResponse)) {
+            if (cumulativeCharCount >= targetCharCount) {
+                return i;
+            }
+            lastSplitPoint = i;
+        }
+        cumulativeCharCount += charCounts[i];
+    }
+    // We found no split points after targetCharCount.
+    // Check if it's safe to compress everything.
+    const lastContent = contents[contents.length - 1];
+    if (lastContent?.role === 'model' &&
+        !lastContent?.parts?.some((part) => part.functionCall)) {
+        return contents.length;
     }
-    return contentLengths.length;
+    // Can't compress everything so just compress at last splitpoint.
+    return lastSplitPoint;
 }
 const MAX_TURNS = 100;
 /**
@@ -62,8 +80,6 @@ const COMPRESSION_PRESERVE_THRESHOLD = 0.3;
 export class GeminiClient {
     config;
     chat;
-    contentGenerator;
-    embeddingModel;
     generateContentConfig = {
         temperature: 0,
         topP: 1,
@@ -71,29 +87,27 @@ export class GeminiClient {
     sessionTurnCount = 0;
     loopDetector;
     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.lastPromptId = this.config.getSessionId();
     }
-    async initialize(contentGeneratorConfig) {
-        this.contentGenerator = await createContentGenerator(contentGeneratorConfig, this.config, this.config.getSessionId());
+    async initialize() {
         this.chat = await this.startChat();
     }
-    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 +119,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() {
@@ -142,6 +140,12 @@ export class GeminiClient {
     async resetChat() {
         this.chat = await this.startChat();
     }
+    getChatRecordingService() {
+        return this.chat?.getChatRecordingService();
+    }
+    getLoopDetectionService() {
+        return this.loopDetector;
+    }
     async addDirectoryContext() {
         if (!this.chat) {
             return;
@@ -153,6 +157,7 @@ export class GeminiClient {
     }
     async startChat(extraHistory) {
         this.forceFullIdeContext = true;
+        this.hasFailedCompressionAttempt = false;
         const envParts = await getEnvironmentContext(this.config);
         const toolRegistry = this.config.getToolRegistry();
         const toolDeclarations = toolRegistry.getFunctionDeclarations();
@@ -171,17 +176,17 @@ export class GeminiClient {
         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(), {
+            const model = this.config.getModel();
+            const config = { ...this.generateContentConfig };
+            if (isThinkingSupported(model)) {
+                config.thinkingConfig = {
+                    includeThoughts: true,
+                    thinkingBudget: DEFAULT_THINKING_MODE,
+                };
+            }
+            return new GeminiChat(this.config, {
                 systemInstruction,
-                ...generateContentConfigWithThinking,
+                ...config,
                 tools,
             }, history);
         }
@@ -191,7 +196,7 @@ export class GeminiClient {
         }
     }
     getIdeContextParts(forceFullContext) {
-        const currentIdeContext = ideContext.getIdeContext();
+        const currentIdeContext = ideContextStore.get();
         if (!currentIdeContext) {
             return { contextParts: [], newIdeContext: undefined };
         }
@@ -329,10 +334,11 @@ export class GeminiClient {
             };
         }
     }
-    async *sendMessageStream(request, signal, prompt_id, turns = MAX_TURNS, originalModel) {
+    async *sendMessageStream(request, signal, prompt_id, turns = MAX_TURNS) {
         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 +351,8 @@ 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) {
+        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,186 +377,63 @@ 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;
+        }
+        const resultStream = turn.run(modelToUse, request, linkedSignal);
         for await (const event of resultStream) {
             if (this.loopDetector.addAndCheck(event)) {
                 yield { type: GeminiEventType.LoopDetected };
+                controller.abort();
                 return turn;
             }
             yield event;
+            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 = {}) {
-        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,
-                    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();
+        let currentAttemptModel = model;
         const configToUse = {
             ...this.generateContentConfig,
             ...generationConfig,
@@ -565,13 +446,22 @@ export class GeminiClient {
                 ...configToUse,
                 systemInstruction,
             };
-            const apiCall = () => this.getContentGenerator().generateContent({
-                model: modelToUse,
-                config: requestConfig,
-                contents,
-            }, this.lastPromptId);
+            const apiCall = () => {
+                const modelToUse = this.config.isInFallbackMode()
+                    ? DEFAULT_GEMINI_FLASH_MODEL
+                    : model;
+                currentAttemptModel = modelToUse;
+                return this.getContentGeneratorOrFail().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,79 +470,70 @@ 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,
             }, 'generateContent-api');
-            throw new Error(`Failed to generate content with model ${modelToUse}: ${getErrorMessage(error)}`);
-        }
-    }
-    async generateEmbedding(texts) {
-        if (!texts || texts.length === 0) {
-            return [];
-        }
-        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}.`);
+            throw new Error(`Failed to generate content with model ${currentAttemptModel}: ${getErrorMessage(error)}`);
         }
-        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) {
+        // 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 configModel = this.config.getModel();
+        let model = configModel === DEFAULT_GEMINI_MODEL_AUTO
+            ? DEFAULT_GEMINI_MODEL
+            : configModel;
+        // Check if the model needs to be a fallback
+        model = getEffectiveModel(this.config.isInFallbackMode(), model);
         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;
+        if (curatedHistory.length === 0 ||
+            (this.hasFailedCompressionAttempt && !force)) {
+            return {
+                originalTokenCount: 0,
+                newTokenCount: 0,
+                compressionStatus: CompressionStatus.NOOP,
+            };
         }
+        const originalTokenCount = uiTelemetryService.getLastPromptTokenCount();
         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;
+                return {
+                    originalTokenCount,
+                    newTokenCount: originalTokenCount,
+                    compressionStatus: CompressionStatus.NOOP,
+                };
             }
         }
-        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>.',
-            },
+        const splitPoint = findCompressSplitPoint(curatedHistory, 1 - COMPRESSION_PRESERVE_THRESHOLD);
+        const historyToCompress = curatedHistory.slice(0, splitPoint);
+        const historyToKeep = curatedHistory.slice(splitPoint);
+        const summaryResponse = await this.config
+            .getContentGenerator()
+            .generateContent({
+            model,
+            contents: [
+                ...historyToCompress,
+                {
+                    role: 'user',
+                    parts: [
+                        {
+                            text: 'First, reason in your scratchpad. Then, generate the <state_snapshot>.',
+                        },
+                    ],
+                },
+            ],
             config: {
                 systemInstruction: { text: getCompressionPrompt() },
             },
         }, prompt_id);
-        this.chat = await this.startChat([
+        const summary = getResponseText(summaryResponse) ?? '';
+        const chat = await this.startChat([
             {
                 role: 'user',
                 parts: [{ text: summary }],
@@ -664,60 +545,33 @@ export class GeminiClient {
             ...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;
-        }
+        // Estimate token count 1 token ≈ 4 characters
+        const newTokenCount = Math.floor(chat
+            .getHistory()
+            .reduce((total, content) => total + JSON.stringify(content).length, 0) /
+            4);
         logChatCompression(this.config, makeChatCompressionEvent({
             tokens_before: originalTokenCount,
             tokens_after: newTokenCount,
         }));
+        if (newTokenCount > originalTokenCount) {
+            this.hasFailedCompressionAttempt = !force && true;
+            return {
+                originalTokenCount,
+                newTokenCount,
+                compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
+            };
+        }
+        else {
+            this.chat = chat; // Chat compression successful, set new state.
+            uiTelemetryService.setLastPromptTokenCount(newTokenCount);
+        }
         return {
             originalTokenCount,
             newTokenCount,
+            compressionStatus: CompressionStatus.COMPRESSED,
         };
     }
-    /**
-     * 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) {
-        // Only handle fallback for OAuth users
-        if (authType !== AuthType.LOGIN_WITH_GOOGLE) {
-            return null;
-        }
-        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) {
-                    this.config.setModel(fallbackModel);
-                    this.config.setFallbackMode(true);
-                    return fallbackModel;
-                }
-                // Check if the model was switched manually in the handler
-                if (this.config.getModel() === fallbackModel) {
-                    return null; // Model was switched but don't continue with current prompt
-                }
-            }
-            catch (error) {
-                console.warn('Flash fallback handler failed:', error);
-            }
-        }
-        return null;
-    }
 }
 export const TEST_ONLY = {
     COMPRESSION_PRESERVE_THRESHOLD,