@office-ai/aioncli-core 0.2.3 → 0.18.4

package/dist/src/core/geminiChat.js

@@ -3,14 +3,31 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-// DISCLAIMER: This is a copied version of https://github.com/googleapis/js-genai/blob/main/src/chats.ts with the intention of working around a key bug
-// where function responses are not treated as "valid" responses: https://b.corp.google.com/issues/420354090
-import { createUserContent, } from '@google/genai';
+import { ThinkingLevel } from '@google/genai';
+import { toParts } from '../code_assist/converter.js';
+import { createUserContent, FinishReason } from '@google/genai';
 import { retryWithBackoff } from '../utils/retry.js';
-import { isFunctionResponse } from '../utils/messageInspectors.js';
-import { AuthType, createContentGenerator } from './contentGenerator.js';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
+import { DEFAULT_GEMINI_MODEL, DEFAULT_THINKING_MODE, PREVIEW_GEMINI_MODEL, getEffectiveModel, isGemini2Model, } from '../config/models.js';
 import { hasCycleInSchema } from '../tools/tools.js';
+import { logContentRetry, logContentRetryFailure, } from '../telemetry/loggers.js';
+import { ChatRecordingService, } from '../services/chatRecordingService.js';
+import { ContentRetryEvent, ContentRetryFailureEvent, } from '../telemetry/types.js';
+import { handleFallback } from '../fallback/handler.js';
+import { isFunctionResponse } from '../utils/messageInspectors.js';
+import { partListUnionToString } from './geminiRequest.js';
+export var StreamEventType;
+(function (StreamEventType) {
+    /** A regular content chunk from the API. */
+    StreamEventType["CHUNK"] = "chunk";
+    /** A signal that a retry is about to happen. The UI should discard any partial
+     * content from the attempt that just failed. */
+    StreamEventType["RETRY"] = "retry";
+})(StreamEventType || (StreamEventType = {}));
+const INVALID_CONTENT_RETRY_OPTIONS = {
+    maxAttempts: 2, // 1 initial call + 1 retry
+    initialDelayMs: 500,
+};
+export const SYNTHETIC_THOUGHT_SIGNATURE = 'skip_thought_signature_validator';
 /**
  * Returns true if the response is valid, false otherwise.
  */
@@ -24,6 +41,16 @@ function isValidResponse(response) {
     }
     return isValidContent(content);
 }
+export function isValidNonThoughtTextPart(part) {
+    return (typeof part.text === 'string' &&
+        !part.thought &&
+        // Technically, the model should never generate parts that have text and
+        //  any of these but we don't trust them so check anyways.
+        !part.functionCall &&
+        !part.functionResponse &&
+        !part.inlineData &&
+        !part.fileData);
+}
 function isValidContent(content) {
     if (content.parts === undefined || content.parts.length === 0) {
         return false;
@@ -84,14 +111,22 @@ function extractCuratedHistory(comprehensiveHistory) {
             if (isValid) {
                 curatedHistory.push(...modelOutput);
             }
-            else {
-                // Remove the last user input when model content is invalid.
-                curatedHistory.pop();
-            }
         }
     }
     return curatedHistory;
 }
+/**
+ * Custom error to signal that a stream completed with invalid content,
+ * which should trigger a retry.
+ */
+export class InvalidStreamError extends Error {
+    type;
+    constructor(message, type) {
+        super(message);
+        this.name = 'InvalidStreamError';
+        this.type = type;
+    }
+}
 /**
  * Chat session that enables sending messages to the model with previous
  * conversation context.
@@ -101,212 +136,26 @@ function extractCuratedHistory(comprehensiveHistory) {
  */
 export class GeminiChat {
     config;
-    contentGenerator;
-    generationConfig;
+    systemInstruction;
+    tools;
     history;
     // A promise to represent the current state of the message being sent to the
     // model.
     sendPromise = Promise.resolve();
-    constructor(config, 
-    // [PATCH:API_KEY_ROTATION] Remove readonly to allow updating contentGenerator for API key rotation
-    contentGenerator, 
-    // [/PATCH:API_KEY_ROTATION]
-    generationConfig = {}, history = []) {
+    chatRecordingService;
+    lastPromptTokenCount;
+    constructor(config, systemInstruction = '', tools = [], history = [], resumedSessionData) {
         this.config = config;
-        this.contentGenerator = contentGenerator;
-        this.generationConfig = generationConfig;
+        this.systemInstruction = systemInstruction;
+        this.tools = tools;
         this.history = history;
         validateHistory(history);
-    }
-    // [PATCH:API_KEY_ROTATION] Update contentGenerator for API key rotation
-    updateContentGenerator(newContentGenerator) {
-        console.log('[API_KEY_ROTATION] Updating content generator in GeminiChat instance');
-        this.contentGenerator = newContentGenerator;
-        console.log('[API_KEY_ROTATION] Content generator updated successfully');
-    }
-    // [/PATCH:API_KEY_ROTATION]
-    /**
-     * [PATCH:API_KEY_ROTATION] Safely format API key for logging (shows first 8 and last 4 chars)
-     */
-    formatApiKeyForLog(apiKey) {
-        if (!apiKey || apiKey.length < 12) {
-            return 'invalid-key';
-        }
-        const start = apiKey.substring(0, 8);
-        const end = apiKey.substring(apiKey.length - 4);
-        const middle = '*'.repeat(Math.min(apiKey.length - 12, 20));
-        return `${start}${middle}${end}`;
-    }
-    /**
-     * [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 (error) {
-                // Fallback handler failed, continue to return null
-            }
-        }
-        return null;
+        this.chatRecordingService = new ChatRecordingService(config);
+        this.chatRecordingService.initialize(resumedSessionData);
+        this.lastPromptTokenCount = Math.ceil(JSON.stringify(this.history).length / 4);
     }
     setSystemInstruction(sysInstr) {
-        this.generationConfig.systemInstruction = sysInstr;
-    }
-    /**
-     * Sends a message to the model and returns the response.
-     *
-     * @remarks
-     * This method will wait for the previous message to be processed before
-     * sending the next message.
-     *
-     * @see {@link Chat#sendMessageStream} for streaming method.
-     * @param params - parameters for sending messages within a chat session.
-     * @returns The model's response.
-     *
-     * @example
-     * ```ts
-     * const chat = ai.chats.create({model: 'gemini-2.0-flash'});
-     * const response = await chat.sendMessage({
-     *   message: 'Why is the sky blue?'
-     * });
-     * console.log(response.text);
-     * ```
-     */
-    async sendMessage(params, prompt_id) {
-        await this.sendPromise;
-        const userContent = createUserContent(params.message);
-        const requestContents = this.getHistory(true).concat(userContent);
-        let response;
-        try {
-            const apiCall = () => {
-                const modelToUse = this.config.getModel() || DEFAULT_GEMINI_FLASH_MODEL;
-                // Prevent Flash model calls immediately after quota error
-                if (this.config.getQuotaErrorOccurred() &&
-                    modelToUse === DEFAULT_GEMINI_FLASH_MODEL) {
-                    throw new Error('Please submit a new query to continue with the Flash model.');
-                }
-                return this.contentGenerator.generateContent({
-                    model: modelToUse,
-                    contents: requestContents,
-                    config: { ...this.generationConfig, ...params.config },
-                }, prompt_id);
-            };
-            response = await retryWithBackoff(apiCall, {
-                shouldRetry: (error) => {
-                    // Check for known error messages and codes.
-                    if (error instanceof Error && error.message) {
-                        if (isSchemaDepthError(error.message))
-                            return false;
-                        if (error.message.includes('429'))
-                            return true;
-                        if (error.message.match(/5\d{2}/))
-                            return true;
-                    }
-                    return false; // Don't retry other errors by default
-                },
-                onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
-                authType: this.config.getContentGeneratorConfig()?.authType,
-            });
-            this.sendPromise = (async () => {
-                const outputContent = response.candidates?.[0]?.content;
-                // Because the AFC input contains the entire curated chat history in
-                // addition to the new user input, we need to truncate the AFC history
-                // to deduplicate the existing chat history.
-                const fullAutomaticFunctionCallingHistory = response.automaticFunctionCallingHistory;
-                const index = this.getHistory(true).length;
-                let automaticFunctionCallingHistory = [];
-                if (fullAutomaticFunctionCallingHistory != null) {
-                    automaticFunctionCallingHistory =
-                        fullAutomaticFunctionCallingHistory.slice(index) ?? [];
-                }
-                const modelOutput = outputContent ? [outputContent] : [];
-                this.recordHistory(userContent, modelOutput, automaticFunctionCallingHistory);
-            })();
-            await this.sendPromise.catch(() => {
-                // Resets sendPromise to avoid subsequent calls failing
-                this.sendPromise = Promise.resolve();
-            });
-            return response;
-        }
-        catch (error) {
-            this.sendPromise = Promise.resolve();
-            throw error;
-        }
+        this.systemInstruction = sysInstr;
     }
     /**
      * Sends a message to the model and returns the response in chunks.
@@ -316,71 +165,172 @@ export class GeminiChat {
      * sending the next message.
      *
      * @see {@link Chat#sendMessage} for non-streaming method.
-     * @param params - parameters for sending the message.
+     * @param modelConfigKey - The key for the model config.
+     * @param message - The list of messages to send.
+     * @param prompt_id - The ID of the prompt.
+     * @param signal - An abort signal for this message.
      * @return The model's response.
      *
      * @example
      * ```ts
      * const chat = ai.chats.create({model: 'gemini-2.0-flash'});
      * const response = await chat.sendMessageStream({
-     *   message: 'Why is the sky blue?'
+     * message: 'Why is the sky blue?'
      * });
      * for await (const chunk of response) {
-     *   console.log(chunk.text);
+     * console.log(chunk.text);
      * }
      * ```
      */
-    async sendMessageStream(params, prompt_id) {
+    async sendMessageStream(modelConfigKey, message, prompt_id, signal) {
         await this.sendPromise;
-        const userContent = createUserContent(params.message);
-        const requestContents = this.getHistory(true).concat(userContent);
-        try {
-            const apiCall = () => {
-                const modelToUse = this.config.getModel();
-                // Prevent Flash model calls immediately after quota error
-                if (this.config.getQuotaErrorOccurred() &&
-                    modelToUse === DEFAULT_GEMINI_FLASH_MODEL) {
-                    throw new Error('Please submit a new query to continue with the Flash model.');
-                }
-                return this.contentGenerator.generateContentStream({
-                    model: modelToUse,
-                    contents: requestContents,
-                    config: { ...this.generationConfig, ...params.config },
-                }, prompt_id);
-            };
-            // Note: Retrying streams can be complex. If generateContentStream itself doesn't handle retries
-            // for transient issues internally before yielding the async generator, this retry will re-initiate
-            // the stream. For simple 429/500 errors on initial call, this is fine.
-            // If errors occur mid-stream, this setup won't resume the stream; it will restart it.
-            const streamResponse = await retryWithBackoff(apiCall, {
-                shouldRetry: (error) => {
-                    // Check for known error messages and codes.
-                    if (error instanceof Error && error.message) {
-                        if (isSchemaDepthError(error.message))
-                            return false;
-                        if (error.message.includes('429'))
-                            return true;
-                        if (error.message.match(/5\d{2}/))
-                            return true;
-                    }
-                    return false; // Don't retry other errors by default
-                },
-                onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
-                authType: this.config.getContentGeneratorConfig()?.authType,
+        // Preview Model Bypass mode for the new request.
+        // This ensures that we attempt to use Preview Model for every new user turn
+        // (unless the "Always" fallback mode is active, which is handled separately).
+        this.config.setPreviewModelBypassMode(false);
+        let streamDoneResolver;
+        const streamDonePromise = new Promise((resolve) => {
+            streamDoneResolver = resolve;
+        });
+        this.sendPromise = streamDonePromise;
+        const userContent = createUserContent(message);
+        const { model, generateContentConfig } = this.config.modelConfigService.getResolvedConfig(modelConfigKey);
+        generateContentConfig.abortSignal = signal;
+        // Record user input - capture complete message with all parts (text, files, images, etc.)
+        // but skip recording function responses (tool call results) as they should be stored in tool call records
+        if (!isFunctionResponse(userContent)) {
+            const userMessage = Array.isArray(message) ? message : [message];
+            const userMessageContent = partListUnionToString(toParts(userMessage));
+            this.chatRecordingService.recordMessage({
+                model,
+                type: 'user',
+                content: userMessageContent,
             });
-            // Resolve the internal tracking of send completion promise - `sendPromise`
-            // for both success and failure response. The actual failure is still
-            // propagated by the `await streamResponse`.
-            this.sendPromise = Promise.resolve(streamResponse)
-                .then(() => undefined)
-                .catch(() => undefined);
-            const result = this.processStreamResponse(streamResponse, userContent);
-            return result;
-        }
-        catch (error) {
-            this.sendPromise = Promise.resolve();
-            throw error;
         }
+        // Add user content to history ONCE before any attempts.
+        this.history.push(userContent);
+        const requestContents = this.getHistory(true);
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const self = this;
+        return (async function* () {
+            try {
+                let lastError = new Error('Request failed after all retries.');
+                let maxAttempts = INVALID_CONTENT_RETRY_OPTIONS.maxAttempts;
+                // If we are in Preview Model Fallback Mode, we want to fail fast (1 attempt)
+                // when probing the Preview Model.
+                if (self.config.isPreviewModelFallbackMode() &&
+                    model === PREVIEW_GEMINI_MODEL) {
+                    maxAttempts = 1;
+                }
+                for (let attempt = 0; attempt < maxAttempts; attempt++) {
+                    try {
+                        if (attempt > 0) {
+                            yield { type: StreamEventType.RETRY };
+                        }
+                        // If this is a retry, set temperature to 1 to encourage different output.
+                        if (attempt > 0) {
+                            generateContentConfig.temperature = 1;
+                        }
+                        const stream = await self.makeApiCallAndProcessStream(model, generateContentConfig, requestContents, prompt_id);
+                        for await (const chunk of stream) {
+                            yield { type: StreamEventType.CHUNK, value: chunk };
+                        }
+                        lastError = null;
+                        break;
+                    }
+                    catch (error) {
+                        lastError = error;
+                        const isContentError = error instanceof InvalidStreamError;
+                        if (isContentError && isGemini2Model(model)) {
+                            // Check if we have more attempts left.
+                            if (attempt < maxAttempts - 1) {
+                                logContentRetry(self.config, new ContentRetryEvent(attempt, error.type, INVALID_CONTENT_RETRY_OPTIONS.initialDelayMs, model));
+                                await new Promise((res) => setTimeout(res, INVALID_CONTENT_RETRY_OPTIONS.initialDelayMs *
+                                    (attempt + 1)));
+                                continue;
+                            }
+                        }
+                        break;
+                    }
+                }
+                if (lastError) {
+                    if (lastError instanceof InvalidStreamError &&
+                        isGemini2Model(model)) {
+                        logContentRetryFailure(self.config, new ContentRetryFailureEvent(maxAttempts, lastError.type, model));
+                    }
+                    throw lastError;
+                }
+                else {
+                    // Preview Model successfully used, disable fallback mode.
+                    // We only do this if we didn't bypass Preview Model (i.e. we actually used it).
+                    if (model === PREVIEW_GEMINI_MODEL &&
+                        !self.config.isPreviewModelBypassMode()) {
+                        self.config.setPreviewModelFallbackMode(false);
+                    }
+                }
+            }
+            finally {
+                streamDoneResolver();
+            }
+        })();
+    }
+    async makeApiCallAndProcessStream(model, generateContentConfig, requestContents, prompt_id) {
+        let effectiveModel = model;
+        const contentsForPreviewModel = this.ensureActiveLoopHasThoughtSignatures(requestContents);
+        const apiCall = () => {
+            let modelToUse = getEffectiveModel(this.config.isInFallbackMode(), model, this.config.getPreviewFeatures());
+            // Preview Model Bypass Logic:
+            // If we are in "Preview Model Bypass Mode" (transient failure), we force downgrade to 2.5 Pro
+            // IF the effective model is currently Preview Model.
+            if (this.config.isPreviewModelBypassMode() &&
+                modelToUse === PREVIEW_GEMINI_MODEL) {
+                modelToUse = DEFAULT_GEMINI_MODEL;
+            }
+            effectiveModel = modelToUse;
+            const config = {
+                ...generateContentConfig,
+                // TODO(12622): Ensure we don't overrwrite these when they are
+                // passed via config.
+                systemInstruction: this.systemInstruction,
+                tools: this.tools,
+            };
+            // TODO(joshualitt): Clean this up with model configs.
+            if (modelToUse.startsWith('gemini-3')) {
+                config.thinkingConfig = {
+                    ...config.thinkingConfig,
+                    thinkingLevel: ThinkingLevel.HIGH,
+                };
+                delete config.thinkingConfig?.thinkingBudget;
+            }
+            else {
+                // The `gemini-3` configs use thinkingLevel, so we have to invert the
+                // change above.
+                config.thinkingConfig = {
+                    ...config.thinkingConfig,
+                    thinkingBudget: DEFAULT_THINKING_MODE,
+                };
+                delete config.thinkingConfig?.thinkingLevel;
+            }
+            return this.config.getContentGenerator().generateContentStream({
+                model: modelToUse,
+                contents: modelToUse === PREVIEW_GEMINI_MODEL
+                    ? contentsForPreviewModel
+                    : requestContents,
+                config,
+            }, prompt_id);
+        };
+        const onPersistent429Callback = async (authType, error) => await handleFallback(this.config, effectiveModel, authType, error);
+        const streamResponse = await retryWithBackoff(apiCall, {
+            onPersistent429: onPersistent429Callback,
+            authType: this.config.getContentGeneratorConfig()?.authType,
+            retryFetchErrors: this.config.getRetryFetchErrors(),
+            signal: generateContentConfig.abortSignal,
+            maxAttempts: this.config.isPreviewModelFallbackMode() &&
+                model === PREVIEW_GEMINI_MODEL
+                ? 1
+                : undefined,
+        });
+        return this.processStreamResponse(model, streamResponse);
     }
     /**
      * Returns the chat history.
@@ -392,7 +342,7 @@ export class GeminiChat {
      * - The `curated history` contains only the valid turns between user and
      * model, which will be included in the subsequent requests sent to the model.
      * - The `comprehensive history` contains all turns, including invalid or
-     *   empty model outputs, providing a complete record of the history.
+     * empty model outputs, providing a complete record of the history.
      *
      * The history is updated after receiving the response from the model,
      * for streaming response, it means receiving the last chunk of the response.
@@ -401,9 +351,9 @@ export class GeminiChat {
      * history`, set the `curated` parameter to `true`.
      *
      * @param curated - whether to return the curated history or the comprehensive
-     *     history.
+     * history.
      * @return History contents alternating between user and model for the entire
-     *     chat session.
+     * chat session.
      */
     getHistory(curated = false) {
         const history = curated
@@ -421,8 +371,6 @@ export class GeminiChat {
     }
     /**
      * Adds a new entry to the chat history.
-     *
-     * @param content - The content to add to the history.
      */
     addHistory(content) {
         this.history.push(content);
@@ -430,8 +378,70 @@ export class GeminiChat {
     setHistory(history) {
         this.history = history;
     }
+    stripThoughtsFromHistory() {
+        this.history = this.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;
+        });
+    }
+    // To ensure our requests validate, the first function call in every model
+    // turn within the active loop must have a `thoughtSignature` property.
+    // If we do not do this, we will get back 400 errors from the API.
+    ensureActiveLoopHasThoughtSignatures(requestContents) {
+        // First, find the start of the active loop by finding the last user turn
+        // with a text message, i.e. that is not a function response.
+        let activeLoopStartIndex = -1;
+        for (let i = requestContents.length - 1; i >= 0; i--) {
+            const content = requestContents[i];
+            if (content.role === 'user' && content.parts?.some((part) => part.text)) {
+                activeLoopStartIndex = i;
+                break;
+            }
+        }
+        if (activeLoopStartIndex === -1) {
+            return requestContents;
+        }
+        // Iterate through every message in the active loop, ensuring that the first
+        // function call in each message's list of parts has a valid
+        // thoughtSignature property. If it does not we replace the function call
+        // with a copy that uses the synthetic thought signature.
+        const newContents = requestContents.slice(); // Shallow copy the array
+        for (let i = activeLoopStartIndex; i < newContents.length; i++) {
+            const content = newContents[i];
+            if (content.role === 'model' && content.parts) {
+                const newParts = content.parts.slice();
+                for (let j = 0; j < newParts.length; j++) {
+                    const part = newParts[j];
+                    if (part.functionCall) {
+                        if (!part.thoughtSignature) {
+                            newParts[j] = {
+                                ...part,
+                                thoughtSignature: SYNTHETIC_THOUGHT_SIGNATURE,
+                            };
+                            newContents[i] = {
+                                ...content,
+                                parts: newParts,
+                            };
+                        }
+                        break; // Only consider the first function call
+                    }
+                }
+            }
+        }
+        return newContents;
+    }
     setTools(tools) {
-        this.generationConfig.tools = tools;
+        this.tools = tools;
     }
     async maybeIncludeSchemaDepthContext(error) {
         // Check for potentially problematic cyclic tools with cyclic schemas
@@ -455,122 +465,134 @@ export class GeminiChat {
             }
         }
     }
-    async *processStreamResponse(streamResponse, inputContent) {
-        const outputContent = [];
-        const chunks = [];
-        let errorOccurred = false;
-        try {
-            for await (const chunk of streamResponse) {
-                if (isValidResponse(chunk)) {
-                    chunks.push(chunk);
-                    const content = chunk.candidates?.[0]?.content;
-                    if (content !== undefined) {
-                        if (this.isThoughtContent(content)) {
-                            yield chunk;
-                            continue;
-                        }
-                        outputContent.push(content);
+    async *processStreamResponse(model, streamResponse) {
+        const modelResponseParts = [];
+        let hasToolCall = false;
+        let finishReason;
+        for await (const chunk of streamResponse) {
+            const candidateWithReason = chunk?.candidates?.find((candidate) => candidate.finishReason);
+            if (candidateWithReason) {
+                finishReason = candidateWithReason.finishReason;
+            }
+            if (isValidResponse(chunk)) {
+                const content = chunk.candidates?.[0]?.content;
+                if (content?.parts) {
+                    if (content.parts.some((part) => part.thought)) {
+                        // Record thoughts
+                        this.recordThoughtFromContent(content);
                     }
+                    if (content.parts.some((part) => part.functionCall)) {
+                        hasToolCall = true;
+                    }
+                    modelResponseParts.push(...content.parts.filter((part) => !part.thought));
                 }
-                yield chunk;
             }
-        }
-        catch (error) {
-            errorOccurred = true;
-            throw error;
-        }
-        if (!errorOccurred) {
-            const allParts = [];
-            for (const content of outputContent) {
-                if (content.parts) {
-                    allParts.push(...content.parts);
+            // Record token usage if this chunk has usageMetadata
+            if (chunk.usageMetadata) {
+                this.chatRecordingService.recordMessageTokens(chunk.usageMetadata);
+                if (chunk.usageMetadata.promptTokenCount !== undefined) {
+                    this.lastPromptTokenCount = chunk.usageMetadata.promptTokenCount;
                 }
             }
-        }
-        this.recordHistory(inputContent, outputContent);
-    }
-    recordHistory(userInput, modelOutput, automaticFunctionCallingHistory) {
-        const nonThoughtModelOutput = modelOutput.filter((content) => !this.isThoughtContent(content));
-        let outputContents = [];
-        if (nonThoughtModelOutput.length > 0 &&
-            nonThoughtModelOutput.every((content) => content.role !== undefined)) {
-            outputContents = nonThoughtModelOutput;
-        }
-        else if (nonThoughtModelOutput.length === 0 && modelOutput.length > 0) {
-            // This case handles when the model returns only a thought.
-            // We don't want to add an empty model response in this case.
-        }
-        else {
-            // When not a function response appends an empty content when model returns empty response, so that the
-            // history is always alternating between user and model.
-            // Workaround for: https://b.corp.google.com/issues/420354090
-            if (!isFunctionResponse(userInput)) {
-                outputContents.push({
-                    role: 'model',
-                    parts: [],
-                });
+            yield chunk; // Yield every chunk to the UI immediately.
+        }
+        // String thoughts and consolidate text parts.
+        const consolidatedParts = [];
+        for (const part of modelResponseParts) {
+            const lastPart = consolidatedParts[consolidatedParts.length - 1];
+            if (lastPart?.text &&
+                isValidNonThoughtTextPart(lastPart) &&
+                isValidNonThoughtTextPart(part)) {
+                lastPart.text += part.text;
+            }
+            else {
+                consolidatedParts.push(part);
             }
         }
-        if (automaticFunctionCallingHistory &&
-            automaticFunctionCallingHistory.length > 0) {
-            this.history.push(...extractCuratedHistory(automaticFunctionCallingHistory));
-        }
-        else {
-            this.history.push(userInput);
+        const responseText = consolidatedParts
+            .filter((part) => part.text)
+            .map((part) => part.text)
+            .join('')
+            .trim();
+        // Record model response text from the collected parts
+        if (responseText) {
+            this.chatRecordingService.recordMessage({
+                model,
+                type: 'gemini',
+                content: responseText,
+            });
         }
-        // Consolidate adjacent model roles in outputContents
-        const consolidatedOutputContents = [];
-        for (const content of outputContents) {
-            if (this.isThoughtContent(content)) {
-                continue;
+        // Stream validation logic: A stream is considered successful if:
+        // 1. There's a tool call OR
+        // 2. A not MALFORMED_FUNCTION_CALL finish reason and a non-mepty resp
+        //
+        // We throw an error only when there's no tool call AND:
+        // - No finish reason, OR
+        // - MALFORMED_FUNCTION_CALL finish reason OR
+        // - Empty response text (e.g., only thoughts with no actual content)
+        if (!hasToolCall) {
+            if (!finishReason) {
+                throw new InvalidStreamError('Model stream ended without a finish reason.', 'NO_FINISH_REASON');
             }
-            const lastContent = consolidatedOutputContents[consolidatedOutputContents.length - 1];
-            if (this.isTextContent(lastContent) && this.isTextContent(content)) {
-                // If both current and last are text, combine their text into the lastContent's first part
-                // and append any other parts from the current content.
-                lastContent.parts[0].text += content.parts[0].text || '';
-                if (content.parts.length > 1) {
-                    lastContent.parts.push(...content.parts.slice(1));
-                }
-            }
-            else {
-                consolidatedOutputContents.push(content);
+            if (finishReason === FinishReason.MALFORMED_FUNCTION_CALL) {
+                throw new InvalidStreamError('Model stream ended with malformed function call.', 'MALFORMED_FUNCTION_CALL');
             }
-        }
-        if (consolidatedOutputContents.length > 0) {
-            const lastHistoryEntry = this.history[this.history.length - 1];
-            const canMergeWithLastHistory = !automaticFunctionCallingHistory ||
-                automaticFunctionCallingHistory.length === 0;
-            if (canMergeWithLastHistory &&
-                this.isTextContent(lastHistoryEntry) &&
-                this.isTextContent(consolidatedOutputContents[0])) {
-                // If both current and last are text, combine their text into the lastHistoryEntry's first part
-                // and append any other parts from the current content.
-                lastHistoryEntry.parts[0].text +=
-                    consolidatedOutputContents[0].parts[0].text || '';
-                if (consolidatedOutputContents[0].parts.length > 1) {
-                    lastHistoryEntry.parts.push(...consolidatedOutputContents[0].parts.slice(1));
-                }
-                consolidatedOutputContents.shift(); // Remove the first element as it's merged
+            if (!responseText) {
+                throw new InvalidStreamError('Model stream ended with empty response text.', 'NO_RESPONSE_TEXT');
             }
-            this.history.push(...consolidatedOutputContents);
         }
+        this.history.push({ role: 'model', parts: consolidatedParts });
+    }
+    getLastPromptTokenCount() {
+        return this.lastPromptTokenCount;
+    }
+    /**
+     * Gets the chat recording service instance.
+     */
+    getChatRecordingService() {
+        return this.chatRecordingService;
     }
-    isTextContent(content) {
-        return !!(content &&
-            content.role === 'model' &&
-            content.parts &&
-            content.parts.length > 0 &&
-            typeof content.parts[0].text === 'string' &&
-            content.parts[0].text !== '');
+    /**
+     * Records completed tool calls with full metadata.
+     * This is called by external components when tool calls complete, before sending responses to Gemini.
+     */
+    recordCompletedToolCalls(model, toolCalls) {
+        const toolCallRecords = toolCalls.map((call) => {
+            const resultDisplayRaw = call.response?.resultDisplay;
+            const resultDisplay = typeof resultDisplayRaw === 'string' ? resultDisplayRaw : undefined;
+            return {
+                id: call.request.callId,
+                name: call.request.name,
+                args: call.request.args,
+                result: call.response?.responseParts || null,
+                status: call.status,
+                timestamp: new Date().toISOString(),
+                resultDisplay,
+            };
+        });
+        this.chatRecordingService.recordToolCalls(model, toolCallRecords);
     }
-    isThoughtContent(content) {
-        return !!(content &&
-            content.role === 'model' &&
-            content.parts &&
-            content.parts.length > 0 &&
-            typeof content.parts[0].thought === 'boolean' &&
-            content.parts[0].thought === true);
+    /**
+     * Extracts and records thought from thought content.
+     */
+    recordThoughtFromContent(content) {
+        if (!content.parts || content.parts.length === 0) {
+            return;
+        }
+        const thoughtPart = content.parts[0];
+        if (thoughtPart.text) {
+            // Extract subject and description using the same logic as turn.ts
+            const rawText = thoughtPart.text;
+            const subjectStringMatches = rawText.match(/\*\*(.*?)\*\*/s);
+            const subject = subjectStringMatches
+                ? subjectStringMatches[1].trim()
+                : '';
+            const description = rawText.replace(/\*\*(.*?)\*\*/s, '').trim();
+            this.chatRecordingService.recordThought({
+                subject,
+                description,
+            });
+        }
     }
 }
 /** Visible for Testing */