++

packages/core/src/config/config.ts

@@ -495,6 +495,7 @@ export interface ConfigParameters {
   mcpEnablementCallbacks?: McpEnablementCallbacks;
   userMemory?: string | HierarchicalMemory;
   geminiMdFileCount?: number;
+  contentGenerator?: ContentGenerator;
   geminiMdFilePaths?: string[];
   approvalMode?: ApprovalMode;
   showMemoryUsage?: boolean;
@@ -566,7 +567,7 @@ export interface ConfigParameters {
   maxAttempts?: number;
   enableShellOutputEfficiency?: boolean;
   shellToolInactivityTimeout?: number;
-  fakeResponses?: string;
+  fakeResponses?: string | any[];
   recordResponses?: string;
   ptyInfo?: string;
   disableYoloMode?: boolean;
@@ -625,6 +626,7 @@ export class Config implements McpContext, AgentLoopContext {
   private trackerService?: TrackerService;
   private contentGeneratorConfig!: ContentGeneratorConfig;
   private contentGenerator!: ContentGenerator;
+  private _initialContentGenerator?: ContentGenerator;
   readonly modelConfigService: ModelConfigService;
   private readonly embeddingModel: string;
   private readonly sandbox: SandboxConfig | undefined;
@@ -764,7 +766,7 @@ export class Config implements McpContext, AgentLoopContext {
   private readonly maxAttempts: number;
   private readonly enableShellOutputEfficiency: boolean;
   private readonly shellToolInactivityTimeout: number;
-  readonly fakeResponses?: string;
+  readonly fakeResponses?: string | any[];
   readonly recordResponses?: string;
   private readonly disableYoloMode: boolean;
   private readonly rawOutput: boolean;
@@ -829,6 +831,7 @@ export class Config implements McpContext, AgentLoopContext {
     this.pendingIncludeDirectories = params.includeDirectories ?? [];
     this.debugMode = params.debugMode;
     this.question = params.question;
+    this._initialContentGenerator = params.contentGenerator;
 
     this.coreTools = params.coreTools;
     this.allowedTools = params.allowedTools;
@@ -1253,11 +1256,17 @@ export class Config implements McpContext, AgentLoopContext {
       baseUrl,
       customHeaders,
     );
-    this.contentGenerator = await createContentGenerator(
-      newContentGeneratorConfig,
-      this,
-      this.getSessionId(),
-    );
+    if (this._initialContentGenerator) {
+      this.contentGenerator = this._initialContentGenerator;
+      // We only use it once, on first initialization. Future refreshes will create real ones
+      // unless we want it to persist forever, but usually AppRig manages this.
+    } else {
+      this.contentGenerator = await createContentGenerator(
+        newContentGeneratorConfig,
+        this,
+        this.getSessionId(),
+      );
+    }
     // Only assign to instance properties after successful initialization
     this.contentGeneratorConfig = newContentGeneratorConfig;
 

packages/core/src/core/contentGenerator.ts

@@ -21,9 +21,10 @@ import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
 import { LoggingContentGenerator } from './loggingContentGenerator.js';
 import { InstallationManager } from '../utils/installationManager.js';
 import { FakeContentGenerator } from './fakeContentGenerator.js';
+import { FallbackContentGenerator } from './fallbackContentGenerator.js';
 import { parseCustomHeaders } from '../utils/customHeaderUtils.js';
 import { RecordingContentGenerator } from './recordingContentGenerator.js';
-import { getVersion, resolveModel } from '../../index.js';
+import { debugLogger, getVersion, resolveModel } from '../../index.js';
 import type { LlmRole } from '../telemetry/llmRole.js';
 
 /**
@@ -160,12 +161,6 @@ export async function createContentGenerator(
   sessionId?: string,
 ): Promise<ContentGenerator> {
   const generator = await (async () => {
-    if (gcConfig.fakeResponses) {
-      const fakeGenerator = await FakeContentGenerator.fromFile(
-        gcConfig.fakeResponses,
-      );
-      return new LoggingContentGenerator(fakeGenerator, gcConfig);
-    }
     const version = await getVersion();
     const model = resolveModel(
       gcConfig.getModel(),
@@ -194,23 +189,21 @@ export async function createContentGenerator(
     ) {
       baseHeaders['Authorization'] = `Bearer ${config.apiKey}`;
     }
+
+    let realGenerator: ContentGenerator;
+
     if (
       config.authType === AuthType.LOGIN_WITH_GOOGLE ||
       config.authType === AuthType.COMPUTE_ADC
     ) {
       const httpOptions = { headers: baseHeaders };
-      return new LoggingContentGenerator(
-        await createCodeAssistContentGenerator(
-          httpOptions,
-          config.authType,
-          gcConfig,
-          sessionId,
-        ),
+      realGenerator = await createCodeAssistContentGenerator(
+        httpOptions,
+        config.authType,
         gcConfig,
+        sessionId,
       );
-    }
-
-    if (
+    } else if (
       config.authType === AuthType.USE_GEMINI ||
       config.authType === AuthType.USE_VERTEX_AI ||
       config.authType === AuthType.GATEWAY
@@ -242,11 +235,28 @@ export async function createContentGenerator(
         httpOptions,
         ...(apiVersionEnv && { apiVersion: apiVersionEnv }),
       });
-      return new LoggingContentGenerator(googleGenAI.models, gcConfig);
+      realGenerator = googleGenAI.models;
+    } else {
+      throw new Error(
+        `Error creating contentGenerator: Unsupported authType: ${config.authType}`,
+      );
     }
-    throw new Error(
-      `Error creating contentGenerator: Unsupported authType: ${config.authType}`,
-    );
+
+    let targetGenerator = realGenerator;
+
+    if (gcConfig.fakeResponses) {
+      if (Array.isArray(gcConfig.fakeResponses)) {
+        debugLogger.log(`[createContentGenerator] Instantiating FakeContentGenerator with ${gcConfig.fakeResponses.length} in-memory mock responses.`);
+        const fakeGen = new FakeContentGenerator(gcConfig.fakeResponses);
+        targetGenerator = new FallbackContentGenerator(fakeGen, realGenerator);
+      } else {
+        debugLogger.log(`[createContentGenerator] Instantiating FakeContentGenerator from file: ${gcConfig.fakeResponses}`);
+        const fakeGen = await FakeContentGenerator.fromFile(gcConfig.fakeResponses);
+        targetGenerator = new FallbackContentGenerator(fakeGen, realGenerator);
+      }
+    }
+
+    return new LoggingContentGenerator(targetGenerator, gcConfig);
   })();
 
   if (gcConfig.recordResponses) {

packages/core/src/core/fakeContentGenerator.ts

@@ -18,6 +18,16 @@ import type { UserTierId, GeminiUserTier } from '../code_assist/types.js';
 import { safeJsonStringify } from '../utils/safeJsonStringify.js';
 import type { LlmRole } from '../telemetry/types.js';
 
+export class MockExhaustedError extends Error {
+  constructor(method: string, request?: unknown) {
+    super(
+      `No more mock responses for ${method}, got request:\n` +
+        safeJsonStringify(request),
+    );
+    this.name = 'MockExhaustedError';
+  }
+}
+
 export type FakeResponse =
   | {
       method: 'generateContent';
@@ -53,7 +63,9 @@ export class FakeContentGenerator implements ContentGenerator {
     return this.sentRequests;
   }
 
-  static async fromFile(filePath: string): Promise<FakeContentGenerator> {
+  static async fromFile(
+    filePath: string,
+  ): Promise<FakeContentGenerator> {
     const fileContent = await promises.readFile(filePath, 'utf-8');
     const responses = fileContent
       .split('\n')
@@ -67,13 +79,14 @@ export class FakeContentGenerator implements ContentGenerator {
     M extends FakeResponse['method'],
     R = Extract<FakeResponse, { method: M }>['response'],
   >(method: M, request: unknown): R {
-    const response = this.responses[this.callCounter++];
+    const response = this.responses[this.callCounter];
     if (!response) {
-      throw new Error(
-        `No more mock responses for ${method}, got request:\n` +
-          safeJsonStringify(request),
-      );
+      throw new MockExhaustedError(method, request);
     }
+
+    // We only increment the counter if we actually consume a mock response
+    this.callCounter++;
+
     if (response.method !== method) {
       throw new Error(
         `Unexpected response type, next response was for ${response.method} but expected ${method}`,
@@ -85,28 +98,29 @@ export class FakeContentGenerator implements ContentGenerator {
 
   async generateContent(
     request: GenerateContentParameters,
-    _userPromptId: string,
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    userPromptId: string,
     role: LlmRole,
   ): Promise<GenerateContentResponse> {
     this.sentRequests.push(request);
+    const next = this.getNextResponse('generateContent', request);
     // eslint-disable-next-line @typescript-eslint/no-unsafe-return
-    return Object.setPrototypeOf(
-      this.getNextResponse('generateContent', request),
-      GenerateContentResponse.prototype,
-    );
+    return Object.setPrototypeOf(next, GenerateContentResponse.prototype);
   }
 
   async generateContentStream(
     request: GenerateContentParameters,
-    _userPromptId: string,
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    userPromptId: string,
     role: LlmRole,
   ): Promise<AsyncGenerator<GenerateContentResponse>> {
     this.sentRequests.push(request);
     const responses = this.getNextResponse('generateContentStream', request);
+    
     async function* stream() {
-      for (const response of responses) {
+      // Add a tiny delay to ensure React has time to render the 'Responding'
+      // state. If the mock stream finishes synchronously, AppRig's
+      // awaitingResponse flag may never be cleared, causing the rig to hang.
+      await new Promise((resolve) => setTimeout(resolve, 5));
+      for (const response of responses!) {
         yield Object.setPrototypeOf(
           response,
           GenerateContentResponse.prototype,
@@ -119,16 +133,15 @@ export class FakeContentGenerator implements ContentGenerator {
   async countTokens(
     request: CountTokensParameters,
   ): Promise<CountTokensResponse> {
-    return this.getNextResponse('countTokens', request);
+    const next = this.getNextResponse('countTokens', request);
+    return next;
   }
 
   async embedContent(
     request: EmbedContentParameters,
   ): Promise<EmbedContentResponse> {
+    const next = this.getNextResponse('embedContent', request);
     // eslint-disable-next-line @typescript-eslint/no-unsafe-return
-    return Object.setPrototypeOf(
-      this.getNextResponse('embedContent', request),
-      EmbedContentResponse.prototype,
-    );
+    return Object.setPrototypeOf(next, EmbedContentResponse.prototype);
   }
 }

packages/core/src/core/fallbackContentGenerator.ts

@@ -0,0 +1,97 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { ContentGenerator } from './contentGenerator.js';
+import type { GenerateContentParameters, GenerateContentResponse, CountTokensParameters, CountTokensResponse, EmbedContentParameters, EmbedContentResponse } from '@google/genai';
+import type { LlmRole } from '../telemetry/types.js';
+import { debugLogger } from '../utils/debugLogger.js';
+import { MockExhaustedError } from './fakeContentGenerator.js';
+
+/**
+ * A ContentGenerator that attempts to use a primary generator,
+ * and falls back to a secondary generator if the primary throws MockExhaustedError.
+ */
+export class FallbackContentGenerator implements ContentGenerator {
+  get userTier() { return this.primary.userTier; }
+  get userTierName() { return this.primary.userTierName; }
+  get paidTier() { return this.primary.paidTier; }
+
+  constructor(
+    private readonly primary: ContentGenerator,
+    private readonly fallback: ContentGenerator,
+    private readonly onFallback?: (method: string) => void,
+  ) {}
+
+  async generateContent(
+    request: GenerateContentParameters,
+    userPromptId: string,
+    role: LlmRole,
+  ): Promise<GenerateContentResponse> {
+    try {
+      return await this.primary.generateContent(request, userPromptId, role);
+    } catch (error) {
+      if (error instanceof MockExhaustedError) {
+        debugLogger.log(`[FallbackContentGenerator] Exhausted primary generator for generateContent. Falling back.`);
+        this.onFallback?.('generateContent');
+        return this.fallback.generateContent(request, userPromptId, role);
+      }
+      throw error;
+    }
+  }
+
+  async generateContentStream(
+    request: GenerateContentParameters,
+    userPromptId: string,
+    role: LlmRole,
+  ): Promise<AsyncGenerator<GenerateContentResponse>> {
+    try {
+      return await this.primary.generateContentStream(request, userPromptId, role);
+    } catch (error) {
+      if (error instanceof MockExhaustedError) {
+        debugLogger.log(`[FallbackContentGenerator] Exhausted primary generator for generateContentStream. Falling back.`);
+        this.onFallback?.('generateContentStream');
+        return this.fallback.generateContentStream(request, userPromptId, role);
+      }
+      throw error;
+    }
+  }
+
+  async countTokens(
+    request: CountTokensParameters,
+  ): Promise<CountTokensResponse> {
+    try {
+      if (!this.primary.countTokens) {
+         throw new MockExhaustedError('countTokens');
+      }
+      return await this.primary.countTokens(request);
+    } catch (error) {
+      if (error instanceof MockExhaustedError && this.fallback.countTokens) {
+        debugLogger.log(`[FallbackContentGenerator] Exhausted primary generator for countTokens. Falling back.`);
+        this.onFallback?.('countTokens');
+        return this.fallback.countTokens(request);
+      }
+      throw error;
+    }
+  }
+
+  async embedContent(
+    request: EmbedContentParameters,
+  ): Promise<EmbedContentResponse> {
+    try {
+      if (!this.primary.embedContent) {
+         throw new MockExhaustedError('embedContent');
+      }
+      return await this.primary.embedContent(request);
+    } catch (error) {
+      if (error instanceof MockExhaustedError && this.fallback.embedContent) {
+        debugLogger.log(`[FallbackContentGenerator] Exhausted primary generator for embedContent. Falling back.`);
+        this.onFallback?.('embedContent');
+        return this.fallback.embedContent(request);
+      }
+      throw error;
+    }
+  }
+}

packages/core/src/core/scriptUtils.ts

@@ -0,0 +1,57 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { GenerateContentResponse } from '@google/genai';
+import type { FakeResponse } from './fakeContentGenerator.js';
+
+export type FakeRequest = { method: 'userText'; text: string };
+export type ScriptItem = FakeResponse | FakeRequest;
+
+export function mockGenerateContentStreamText(text: string): FakeResponse {
+  return {
+    method: 'generateContentStream',
+    response: [
+      {
+        candidates: [{ content: { parts: [{ text }] }, finishReason: 'STOP' }],
+      },
+    ] as GenerateContentResponse[],
+  };
+}
+
+export function mockGenerateContentText(text: string): FakeResponse {
+  return {
+    method: 'generateContent',
+    response: {
+      candidates: [{ content: { parts: [{ text }] }, finishReason: 'STOP' }],
+    } as GenerateContentResponse,
+  };
+}
+
+export function userText(text: string): FakeRequest {
+  return { method: 'userText', text };
+}
+
+export function isFakeResponse(item: ScriptItem): item is FakeResponse {
+  return item.method !== 'userText';
+}
+
+export function isFakeRequest(item: ScriptItem): item is FakeRequest {
+  return item.method === 'userText';
+}
+
+/**
+ * Extracts all FakeRequests from a script array and maps them to their string text.
+ */
+export function extractUserPrompts(script: ScriptItem[]): string[] {
+  return script.filter(isFakeRequest).map((req) => req.text);
+}
+
+/**
+ * Extracts all FakeResponses from a script array.
+ */
+export function extractFakeResponses(script: ScriptItem[]): FakeResponse[] {
+  return script.filter(isFakeResponse);
+}

packages/core/src/index.ts

@@ -35,6 +35,9 @@ export * from './commands/types.js';
 export * from './core/baseLlmClient.js';
 export * from './core/client.js';
 export * from './core/contentGenerator.js';
+export * from './core/fakeContentGenerator.js';
+export * from './core/fallbackContentGenerator.js';
+export * from './core/scriptUtils.js';
 export * from './core/loggingContentGenerator.js';
 export * from './core/geminiChat.js';
 export * from './core/logger.js';