@machina.ai/cell-cli-core 1.0.21-rc4 → 1.4.0-rc1

package/dist/src/core/geminiChat.js

@@ -3,14 +3,29 @@
  * 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 { toParts } from '../code_assist/converter.js';
+import { createUserContent } from '@google/genai';
 import { retryWithBackoff } from '../utils/retry.js';
-import { isFunctionResponse } from '../utils/messageInspectors.js';
 import { AuthType } from './contentGenerator.js';
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { hasCycleInSchema } from '../tools/tools.js';
+import { logContentRetry, logContentRetryFailure, logInvalidChunk, } from '../telemetry/loggers.js';
+import { ChatRecordingService } from '../services/chatRecordingService.js';
+import { ContentRetryEvent, ContentRetryFailureEvent, InvalidChunkEvent, } from '../telemetry/types.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: 3, // 1 initial call + 2 retries
+    initialDelayMs: 500,
+};
 /**
  * Returns true if the response is valid, false otherwise.
  */
@@ -84,14 +99,20 @@ 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 without valid content,
+ * which should trigger a retry.
+ */
+export class EmptyStreamError extends Error {
+    constructor(message) {
+        super(message);
+        this.name = 'EmptyStreamError';
+    }
+}
 /**
  * Chat session that enables sending messages to the model with previous
  * conversation context.
@@ -107,12 +128,15 @@ export class GeminiChat {
     // A promise to represent the current state of the message being sent to the
     // model.
     sendPromise = Promise.resolve();
+    chatRecordingService;
     constructor(config, contentGenerator, generationConfig = {}, history = []) {
         this.config = config;
         this.contentGenerator = contentGenerator;
         this.generationConfig = generationConfig;
         this.history = history;
         validateHistory(history);
+        this.chatRecordingService = new ChatRecordingService(config);
+        this.chatRecordingService.initialize();
     }
     /**
      * Handles falling back to Flash model when persistent 429 errors occur for OAuth users.
@@ -168,7 +192,7 @@ export class GeminiChat {
      * ```ts
      * const chat = ai.chats.create({model: 'gemini-2.0-flash'});
      * const response = await chat.sendMessage({
-     *   message: 'Why is the sky blue?'
+     * message: 'Why is the sky blue?'
      * });
      * console.log(response.text);
      * ```
@@ -176,6 +200,17 @@ export class GeminiChat {
     async sendMessage(params, prompt_id) {
         await this.sendPromise;
         const userContent = createUserContent(params.message);
+        // 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(params.message)
+                ? params.message
+                : [params.message];
+            this.chatRecordingService.recordMessage({
+                type: 'user',
+                content: userMessage,
+            });
+        }
         const requestContents = this.getHistory(true).concat(userContent);
         let response;
         try {
@@ -210,6 +245,7 @@ export class GeminiChat {
             });
             this.sendPromise = (async () => {
                 const outputContent = response.candidates?.[0]?.content;
+                const modelOutput = outputContent ? [outputContent] : [];
                 // 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.
@@ -220,12 +256,13 @@ export class GeminiChat {
                     automaticFunctionCallingHistory =
                         fullAutomaticFunctionCallingHistory.slice(index) ?? [];
                 }
-                const modelOutput = outputContent ? [outputContent] : [];
                 this.recordHistory(userContent, modelOutput, automaticFunctionCallingHistory);
             })();
-            await this.sendPromise.catch(() => {
+            await this.sendPromise.catch((error) => {
                 // Resets sendPromise to avoid subsequent calls failing
                 this.sendPromise = Promise.resolve();
+                // Re-throw the error so the caller knows something went wrong.
+                throw error;
             });
             return response;
         }
@@ -249,64 +286,113 @@ export class GeminiChat {
      * ```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) {
         await this.sendPromise;
+        let streamDoneResolver;
+        const streamDonePromise = new Promise((resolve) => {
+            streamDoneResolver = resolve;
+        });
+        this.sendPromise = streamDonePromise;
         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,
+        // 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(params.message)
+                ? params.message
+                : [params.message];
+            const userMessageContent = partListUnionToString(toParts(userMessage));
+            this.chatRecordingService.recordMessage({
+                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.');
+                for (let attempt = 0; attempt < INVALID_CONTENT_RETRY_OPTIONS.maxAttempts; attempt++) {
+                    try {
+                        if (attempt > 0) {
+                            yield { type: StreamEventType.RETRY };
+                        }
+                        const stream = await self.makeApiCallAndProcessStream(requestContents, params, prompt_id, userContent);
+                        for await (const chunk of stream) {
+                            yield { type: StreamEventType.CHUNK, value: chunk };
+                        }
+                        lastError = null;
+                        break;
+                    }
+                    catch (error) {
+                        lastError = error;
+                        const isContentError = error instanceof EmptyStreamError;
+                        if (isContentError) {
+                            // Check if we have more attempts left.
+                            if (attempt < INVALID_CONTENT_RETRY_OPTIONS.maxAttempts - 1) {
+                                logContentRetry(self.config, new ContentRetryEvent(attempt, 'EmptyStreamError', INVALID_CONTENT_RETRY_OPTIONS.initialDelayMs));
+                                await new Promise((res) => setTimeout(res, INVALID_CONTENT_RETRY_OPTIONS.initialDelayMs *
+                                    (attempt + 1)));
+                                continue;
+                            }
+                        }
+                        break;
+                    }
+                }
+                if (lastError) {
+                    if (lastError instanceof EmptyStreamError) {
+                        logContentRetryFailure(self.config, new ContentRetryFailureEvent(INVALID_CONTENT_RETRY_OPTIONS.maxAttempts, 'EmptyStreamError'));
+                    }
+                    // If the stream fails, remove the user message that was added.
+                    if (self.history[self.history.length - 1] === userContent) {
+                        self.history.pop();
+                    }
+                    throw lastError;
+                }
+            }
+            finally {
+                streamDoneResolver();
+            }
+        })();
+    }
+    async makeApiCallAndProcessStream(requestContents, params, prompt_id, userContent) {
+        const apiCall = () => {
+            const modelToUse = this.config.getModel();
+            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);
+        };
+        const streamResponse = await retryWithBackoff(apiCall, {
+            shouldRetry: (error) => {
+                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;
+            },
+            onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
+            authType: this.config.getContentGeneratorConfig()?.authType,
+        });
+        return this.processStreamResponse(streamResponse, userContent);
     }
     /**
      * Returns the chat history.
@@ -318,7 +404,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.
@@ -327,9 +413,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
@@ -347,8 +433,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);
@@ -364,7 +448,7 @@ export class GeminiChat {
         // and include a recommendation to remove potentially problematic tools.
         if (isSchemaDepthError(error.message) ||
             isInvalidArgumentError(error.message)) {
-            const tools = (await this.config.getToolRegistry()).getAllTools();
+            const tools = this.config.getToolRegistry().getAllTools();
             const cyclicSchemaTools = [];
             for (const tool of tools) {
                 if ((tool.schema.parametersJsonSchema &&
@@ -381,122 +465,211 @@ 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(streamResponse, userInput) {
+        const modelResponseParts = [];
+        let hasReceivedAnyChunk = false;
+        let hasReceivedValidChunk = false;
+        let hasToolCall = false;
+        let lastChunk = null;
+        let lastChunkIsInvalid = false;
+        for await (const chunk of streamResponse) {
+            hasReceivedAnyChunk = true;
+            lastChunk = chunk;
+            if (isValidResponse(chunk)) {
+                hasReceivedValidChunk = true;
+                lastChunkIsInvalid = false;
+                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;
+                    }
+                    // Always add parts - thoughts will be filtered out later in recordHistory
+                    modelResponseParts.push(...content.parts);
+                    if (content.parts.some((part) => part.functionCall)) {
+                        hasToolCall = true;
                     }
                 }
-                yield chunk;
             }
-        }
-        catch (error) {
-            errorOccurred = true;
-            throw error;
-        }
-        if (!errorOccurred) {
-            const allParts = [];
-            for (const content of outputContent) {
-                if (content.parts) {
-                    allParts.push(...content.parts);
-                }
+            else {
+                logInvalidChunk(this.config, new InvalidChunkEvent('Invalid chunk received from stream.'));
+                lastChunkIsInvalid = true;
             }
+            // Record token usage if this chunk has usageMetadata
+            if (chunk.usageMetadata) {
+                this.chatRecordingService.recordMessageTokens(chunk.usageMetadata);
+            }
+            yield chunk; // Yield every chunk to the UI immediately.
         }
-        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;
+        if (!hasReceivedAnyChunk) {
+            throw new EmptyStreamError('Model stream completed without any chunks.');
         }
-        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.
+        const hasFinishReason = lastChunk?.candidates?.some((candidate) => candidate.finishReason);
+        // Stream validation logic: A stream is considered successful if:
+        // 1. There's a tool call (tool calls can end without explicit finish reasons), OR
+        // 2. There's a finish reason AND the last chunk is valid (or we haven't received any valid chunks)
+        //
+        // We throw an error only when there's no tool call AND:
+        // - No finish reason, OR
+        // - Last chunk is invalid after receiving valid content
+        if (!hasToolCall &&
+            (!hasFinishReason || (lastChunkIsInvalid && !hasReceivedValidChunk))) {
+            throw new EmptyStreamError('Model stream ended with an invalid chunk or missing finish reason.');
         }
-        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: [],
+        // Record model response text from the collected parts
+        if (modelResponseParts.length > 0) {
+            const responseText = modelResponseParts
+                .filter((part) => part.text && !part.thought)
+                .map((part) => part.text)
+                .join('');
+            if (responseText.trim()) {
+                this.chatRecordingService.recordMessage({
+                    type: 'gemini',
+                    content: responseText,
                 });
             }
         }
+        // Bundle all streamed parts into a single Content object
+        const modelOutput = modelResponseParts.length > 0
+            ? [{ role: 'model', parts: modelResponseParts }]
+            : [];
+        // Pass the raw, bundled data to the new, robust recordHistory
+        this.recordHistory(userInput, modelOutput);
+    }
+    recordHistory(userInput, modelOutput, automaticFunctionCallingHistory) {
+        // Part 1: Handle the user's turn.
         if (automaticFunctionCallingHistory &&
             automaticFunctionCallingHistory.length > 0) {
             this.history.push(...extractCuratedHistory(automaticFunctionCallingHistory));
         }
         else {
-            this.history.push(userInput);
+            if (this.history.length === 0 ||
+                this.history[this.history.length - 1] !== userInput) {
+                const lastTurn = this.history[this.history.length - 1];
+                // The only time we don't push is if it's the *exact same* object,
+                // which happens in streaming where we add it preemptively.
+                if (lastTurn !== userInput) {
+                    if (lastTurn?.role === 'user') {
+                        // This is an invalid sequence.
+                        throw new Error('Cannot add a user turn after another user turn.');
+                    }
+                    this.history.push(userInput);
+                }
+            }
         }
-        // Consolidate adjacent model roles in outputContents
-        const consolidatedOutputContents = [];
-        for (const content of outputContents) {
-            if (this.isThoughtContent(content)) {
+        // Part 2: Process the model output into a final, consolidated list of turns.
+        const finalModelTurns = [];
+        for (const content of modelOutput) {
+            // A. Preserve malformed content that has no 'parts' array.
+            if (!content.parts) {
+                finalModelTurns.push(content);
                 continue;
             }
-            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));
-                }
+            // B. Filter out 'thought' parts.
+            const visibleParts = content.parts.filter((part) => !part.thought);
+            const newTurn = { ...content, parts: visibleParts };
+            const lastTurnInFinal = finalModelTurns[finalModelTurns.length - 1];
+            // Consolidate this new turn with the PREVIOUS turn if they are adjacent model turns.
+            if (lastTurnInFinal &&
+                lastTurnInFinal.role === 'model' &&
+                newTurn.role === 'model' &&
+                lastTurnInFinal.parts && // SAFETY CHECK: Ensure the destination has a parts array.
+                newTurn.parts) {
+                lastTurnInFinal.parts.push(...newTurn.parts);
             }
             else {
-                consolidatedOutputContents.push(content);
+                finalModelTurns.push(newTurn);
             }
         }
-        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));
+        // Part 3: Add the processed model turns to the history, with one final consolidation pass.
+        if (finalModelTurns.length > 0) {
+            // Re-consolidate parts within any turns that were merged in the previous step.
+            for (const turn of finalModelTurns) {
+                if (turn.parts && turn.parts.length > 1) {
+                    const consolidatedParts = [];
+                    for (const part of turn.parts) {
+                        const lastPart = consolidatedParts[consolidatedParts.length - 1];
+                        if (lastPart &&
+                            // Ensure lastPart is a pure text part
+                            typeof lastPart.text === 'string' &&
+                            !lastPart.functionCall &&
+                            !lastPart.functionResponse &&
+                            !lastPart.inlineData &&
+                            !lastPart.fileData &&
+                            !lastPart.thought &&
+                            // Ensure current part is a pure text part
+                            typeof part.text === 'string' &&
+                            !part.functionCall &&
+                            !part.functionResponse &&
+                            !part.inlineData &&
+                            !part.fileData &&
+                            !part.thought) {
+                            lastPart.text += part.text;
+                        }
+                        else {
+                            consolidatedParts.push({ ...part });
+                        }
+                    }
+                    turn.parts = consolidatedParts;
                 }
-                consolidatedOutputContents.shift(); // Remove the first element as it's merged
             }
-            this.history.push(...consolidatedOutputContents);
+            this.history.push(...finalModelTurns);
+        }
+        else {
+            // If, after all processing, there's NO model output, add the placeholder.
+            this.history.push({ role: 'model', parts: [] });
         }
     }
-    isTextContent(content) {
-        return !!(content &&
-            content.role === 'model' &&
-            content.parts &&
-            content.parts.length > 0 &&
-            typeof content.parts[0].text === 'string' &&
-            content.parts[0].text !== '');
+    /**
+     * Gets the chat recording service instance.
+     */
+    getChatRecordingService() {
+        return this.chatRecordingService;
+    }
+    /**
+     * Records completed tool calls with full metadata.
+     * This is called by external components when tool calls complete, before sending responses to Gemini.
+     */
+    recordCompletedToolCalls(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(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 */