@machina.ai/cell-cli-core 1.4.0-rc2 → 1.6.1-rc1

package/dist/src/core/geminiChat.d.ts

@@ -3,11 +3,9 @@
  * Copyright 2025 Google LLC
  * SPDX-License-Identifier: Apache-2.0
  */
-import type { GenerateContentResponse, Content, GenerateContentConfig, SendMessageParameters, Tool } from '@google/genai';
-import type { ContentGenerator } from './contentGenerator.js';
+import { GenerateContentResponse, type Content, type GenerateContentConfig, type SendMessageParameters, type Part, type Tool } from '@google/genai';
 import type { Config } from '../config/config.js';
 import type { StructuredError } from './turn.js';
-import type { CompletedToolCall } from './coreToolScheduler.js';
 import { ChatRecordingService } from '../services/chatRecordingService.js';
 export declare enum StreamEventType {
     /** A regular content chunk from the API. */
@@ -22,6 +20,7 @@ export type StreamEvent = {
 } | {
     type: StreamEventType.RETRY;
 };
+export declare function isValidNonThoughtTextPart(part: Part): boolean;
 /**
  * Custom error to signal that a stream completed without valid content,
  * which should trigger a retry.
@@ -38,39 +37,12 @@ export declare class EmptyStreamError extends Error {
  */
 export declare class GeminiChat {
     private readonly config;
-    private readonly contentGenerator;
     private readonly generationConfig;
     private history;
     private sendPromise;
     private readonly chatRecordingService;
-    constructor(config: Config, contentGenerator: ContentGenerator, generationConfig?: GenerateContentConfig, history?: Content[]);
-    /**
-     * 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.
-     */
-    private handleFlashFallback;
+    constructor(config: Config, generationConfig?: GenerateContentConfig, history?: Content[]);
     setSystemInstruction(sysInstr: string): void;
-    /**
-     * 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);
-     * ```
-     */
-    sendMessage(params: SendMessageParameters, prompt_id: string): Promise<GenerateContentResponse>;
     /**
      * Sends a message to the model and returns the response in chunks.
      *
@@ -93,7 +65,7 @@ export declare class GeminiChat {
      * }
      * ```
      */
-    sendMessageStream(params: SendMessageParameters, prompt_id: string): Promise<AsyncGenerator<StreamEvent>>;
+    sendMessageStream(model: string, params: SendMessageParameters, prompt_id: string): Promise<AsyncGenerator<StreamEvent>>;
     private makeApiCallAndProcessStream;
     /**
      * Returns the chat history.
@@ -128,23 +100,28 @@ export declare class GeminiChat {
      */
     addHistory(content: Content): void;
     setHistory(history: Content[]): void;
+    stripThoughtsFromHistory(): void;
     setTools(tools: Tool[]): void;
     maybeIncludeSchemaDepthContext(error: StructuredError): Promise<void>;
     private processStreamResponse;
-    private recordHistory;
     /**
      * Gets the chat recording service instance.
      */
     getChatRecordingService(): 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: CompletedToolCall[]): void;
     /**
      * Extracts and records thought from thought content.
      */
     private recordThoughtFromContent;
+    /**
+     * Truncates the chunkStream right before the second function call to a
+     * function that mutates state. This may involve trimming parts from a chunk
+     * as well as omtting some chunks altogether.
+     *
+     * We do this because it improves tool call quality if the model gets
+     * feedback from one mutating function call before it makes the next one.
+     */
+    private stopBeforeSecondMutator;
+    private isMutatorFunctionCall;
 }
 /** Visible for Testing */
 export declare function isSchemaDepthError(errorMessage: string): boolean;

package/dist/src/core/geminiChat.js

@@ -3,15 +3,18 @@
  * 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 { GenerateContentResponse, FinishReason, } from '@google/genai';
 import { toParts } from '../code_assist/converter.js';
 import { createUserContent } from '@google/genai';
 import { retryWithBackoff } from '../utils/retry.js';
-import { AuthType } from './contentGenerator.js';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
-import { hasCycleInSchema } from '../tools/tools.js';
+import { DEFAULT_GEMINI_FLASH_MODEL, getEffectiveModel, } from '../config/models.js';
+import { hasCycleInSchema, MUTATOR_KINDS } 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 { handleFallback } from '../fallback/handler.js';
 import { isFunctionResponse } from '../utils/messageInspectors.js';
 import { partListUnionToString } from './geminiRequest.js';
 export var StreamEventType;
@@ -39,6 +42,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;
@@ -122,155 +135,23 @@ export class EmptyStreamError extends Error {
  */
 export class GeminiChat {
     config;
-    contentGenerator;
     generationConfig;
     history;
     // A promise to represent the current state of the message being sent to the
     // model.
     sendPromise = Promise.resolve();
     chatRecordingService;
-    constructor(config, contentGenerator, generationConfig = {}, history = []) {
+    constructor(config, 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.
-     * 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;
-    }
     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);
-        // 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 {
-            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;
-                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.
-                const fullAutomaticFunctionCallingHistory = response.automaticFunctionCallingHistory;
-                const index = this.getHistory(true).length;
-                let automaticFunctionCallingHistory = [];
-                if (fullAutomaticFunctionCallingHistory != null) {
-                    automaticFunctionCallingHistory =
-                        fullAutomaticFunctionCallingHistory.slice(index) ?? [];
-                }
-                this.recordHistory(userContent, modelOutput, automaticFunctionCallingHistory);
-            })();
-            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;
-        }
-        catch (error) {
-            this.sendPromise = Promise.resolve();
-            throw error;
-        }
-    }
     /**
      * Sends a message to the model and returns the response in chunks.
      *
@@ -293,7 +174,7 @@ export class GeminiChat {
      * }
      * ```
      */
-    async sendMessageStream(params, prompt_id) {
+    async sendMessageStream(model, params, prompt_id) {
         await this.sendPromise;
         let streamDoneResolver;
         const streamDonePromise = new Promise((resolve) => {
@@ -309,6 +190,7 @@ export class GeminiChat {
                 : [params.message];
             const userMessageContent = partListUnionToString(toParts(userMessage));
             this.chatRecordingService.recordMessage({
+                model,
                 type: 'user',
                 content: userMessageContent,
             });
@@ -326,7 +208,7 @@ export class GeminiChat {
                         if (attempt > 0) {
                             yield { type: StreamEventType.RETRY };
                         }
-                        const stream = await self.makeApiCallAndProcessStream(requestContents, params, prompt_id, userContent);
+                        const stream = await self.makeApiCallAndProcessStream(model, requestContents, params, prompt_id);
                         for await (const chunk of stream) {
                             yield { type: StreamEventType.CHUNK, value: chunk };
                         }
@@ -364,19 +246,20 @@ export class GeminiChat {
             }
         })();
     }
-    async makeApiCallAndProcessStream(requestContents, params, prompt_id, userContent) {
+    async makeApiCallAndProcessStream(model, requestContents, params, prompt_id) {
         const apiCall = () => {
-            const modelToUse = this.config.getModel();
+            const modelToUse = getEffectiveModel(this.config.isInFallbackMode(), model);
             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({
+            return this.config.getContentGenerator().generateContentStream({
                 model: modelToUse,
                 contents: requestContents,
                 config: { ...this.generationConfig, ...params.config },
             }, prompt_id);
         };
+        const onPersistent429Callback = async (authType, error) => await handleFallback(this.config, model, authType, error);
         const streamResponse = await retryWithBackoff(apiCall, {
             shouldRetry: (error) => {
                 if (error instanceof Error && error.message) {
@@ -389,10 +272,10 @@ export class GeminiChat {
                 }
                 return false;
             },
-            onPersistent429: async (authType, error) => await this.handleFlashFallback(authType, error),
+            onPersistent429: onPersistent429Callback,
             authType: this.config.getContentGeneratorConfig()?.authType,
         });
-        return this.processStreamResponse(streamResponse, userContent);
+        return this.processStreamResponse(model, streamResponse);
     }
     /**
      * Returns the chat history.
@@ -440,6 +323,22 @@ 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;
+        });
+    }
     setTools(tools) {
         this.generationConfig.tools = tools;
     }
@@ -465,14 +364,14 @@ export class GeminiChat {
             }
         }
     }
-    async *processStreamResponse(streamResponse, userInput) {
+    async *processStreamResponse(model, streamResponse) {
         const modelResponseParts = [];
         let hasReceivedAnyChunk = false;
         let hasReceivedValidChunk = false;
         let hasToolCall = false;
         let lastChunk = null;
         let lastChunkIsInvalid = false;
-        for await (const chunk of streamResponse) {
+        for await (const chunk of this.stopBeforeSecondMutator(streamResponse)) {
             hasReceivedAnyChunk = true;
             lastChunk = chunk;
             if (isValidResponse(chunk)) {
@@ -487,11 +386,7 @@ export class GeminiChat {
                     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;
-                    }
+                    modelResponseParts.push(...content.parts.filter((part) => !part.thought));
                 }
             }
             else {
@@ -522,106 +417,31 @@ export class GeminiChat {
         // Record model response text from the collected parts
         if (modelResponseParts.length > 0) {
             const responseText = modelResponseParts
-                .filter((part) => part.text && !part.thought)
+                .filter((part) => part.text)
                 .map((part) => part.text)
                 .join('');
             if (responseText.trim()) {
                 this.chatRecordingService.recordMessage({
+                    model,
                     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 {
-            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);
-                }
-            }
-        }
-        // 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;
-            }
-            // 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);
+        // 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 {
-                finalModelTurns.push(newTurn);
-            }
-        }
-        // 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;
-                }
+                consolidatedParts.push(part);
             }
-            this.history.push(...finalModelTurns);
-        }
-        else {
-            // If, after all processing, there's NO model output, add the placeholder.
-            this.history.push({ role: 'model', parts: [] });
         }
+        this.history.push({ role: 'model', parts: consolidatedParts });
     }
     /**
      * Gets the chat recording service instance.
@@ -629,26 +449,6 @@ export class GeminiChat {
     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);
-    }
     /**
      * Extracts and records thought from thought content.
      */
@@ -671,6 +471,57 @@ export class GeminiChat {
             });
         }
     }
+    /**
+     * Truncates the chunkStream right before the second function call to a
+     * function that mutates state. This may involve trimming parts from a chunk
+     * as well as omtting some chunks altogether.
+     *
+     * We do this because it improves tool call quality if the model gets
+     * feedback from one mutating function call before it makes the next one.
+     */
+    async *stopBeforeSecondMutator(chunkStream) {
+        let foundMutatorFunctionCall = false;
+        for await (const chunk of chunkStream) {
+            const candidate = chunk.candidates?.[0];
+            const content = candidate?.content;
+            if (!candidate || !content?.parts) {
+                yield chunk;
+                continue;
+            }
+            const truncatedParts = [];
+            for (const part of content.parts) {
+                if (this.isMutatorFunctionCall(part)) {
+                    if (foundMutatorFunctionCall) {
+                        // This is the second mutator call.
+                        // Truncate and return immedaitely.
+                        const newChunk = new GenerateContentResponse();
+                        newChunk.candidates = [
+                            {
+                                ...candidate,
+                                content: {
+                                    ...content,
+                                    parts: truncatedParts,
+                                },
+                                finishReason: FinishReason.STOP,
+                            },
+                        ];
+                        yield newChunk;
+                        return;
+                    }
+                    foundMutatorFunctionCall = true;
+                }
+                truncatedParts.push(part);
+            }
+            yield chunk;
+        }
+    }
+    isMutatorFunctionCall(part) {
+        if (!part?.functionCall?.name) {
+            return false;
+        }
+        const tool = this.config.getToolRegistry().getTool(part.functionCall.name);
+        return !!tool && MUTATOR_KINDS.includes(tool.kind);
+    }
 }
 /** Visible for Testing */
 export function isSchemaDepthError(errorMessage) {