From bbff78de672a5ff73314e852a5ca4b41ecfd6248 Mon Sep 17 00:00:00 2001 From: Your Name Date: Wed, 13 May 2026 22:10:26 +0000 Subject: [PATCH] refactor(context): merge non-strict logic back into FakeContentGenerator --- packages/core/src/core/contentGenerator.ts | 4 +- .../core/src/core/fakeContentGenerator.ts | 37 +++++- .../src/core/nonStrictFakeContentGenerator.ts | 111 ------------------ 3 files changed, 36 insertions(+), 116 deletions(-) delete mode 100644 packages/core/src/core/nonStrictFakeContentGenerator.ts diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index a2ac1024e9..9382baac0d 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -23,7 +23,6 @@ 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 { NonStrictFakeContentGenerator } from './nonStrictFakeContentGenerator.js'; import { parseCustomHeaders } from '../utils/customHeaderUtils.js'; import { determineSurface } from '../utils/surface.js'; import { RecordingContentGenerator } from './recordingContentGenerator.js'; @@ -195,8 +194,9 @@ export async function createContentGenerator( ): Promise { const generator = await (async () => { if (gcConfig.fakeResponsesNonStrict) { - const fakeGenerator = await NonStrictFakeContentGenerator.fromFile( + const fakeGenerator = await FakeContentGenerator.fromFile( gcConfig.fakeResponsesNonStrict, + { nonStrict: true }, ); return new LoggingContentGenerator(fakeGenerator, gcConfig); } diff --git a/packages/core/src/core/fakeContentGenerator.ts b/packages/core/src/core/fakeContentGenerator.ts index 39687579e8..0f5a1e35d0 100644 --- a/packages/core/src/core/fakeContentGenerator.ts +++ b/packages/core/src/core/fakeContentGenerator.ts @@ -36,6 +36,18 @@ export type FakeResponse = response: EmbedContentResponse; }; +/** + * Options for the FakeContentGenerator. + */ +export interface FakeContentGeneratorOptions { + /** + * If true, the generator will find the first available response that matches + * the requested method, rather than strictly following the input order. + * Useful for non-deterministic background tasks. + */ + nonStrict?: boolean; +} + // A ContentGenerator that responds with canned responses. // // Typically these would come from a file, provided by the `--fake-responses` @@ -46,22 +58,41 @@ export class FakeContentGenerator implements ContentGenerator { userTierName?: string; paidTier?: GeminiUserTier; - constructor(private readonly responses: FakeResponse[]) {} + constructor( + private readonly responses: FakeResponse[], + private readonly options: FakeContentGeneratorOptions = {}, + ) {} - static async fromFile(filePath: string): Promise { + static async fromFile( + filePath: string, + options: FakeContentGeneratorOptions = {}, + ): Promise { const fileContent = await promises.readFile(filePath, 'utf-8'); const responses = fileContent .split('\n') .filter((line) => line.trim() !== '') // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion .map((line) => JSON.parse(line) as FakeResponse); - return new FakeContentGenerator(responses); + return new FakeContentGenerator(responses, options); } private getNextResponse< M extends FakeResponse['method'], R = Extract['response'], >(method: M, request: unknown): R { + if (this.options.nonStrict) { + const index = this.responses.findIndex((r) => r.method === method); + if (index === -1) { + throw new Error( + `No more mock responses for ${method}, got request:\n` + + safeJsonStringify(request), + ); + } + const response = this.responses.splice(index, 1)[0]; + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return response.response as R; + } + const response = this.responses[this.callCounter++]; if (!response) { throw new Error( diff --git a/packages/core/src/core/nonStrictFakeContentGenerator.ts b/packages/core/src/core/nonStrictFakeContentGenerator.ts deleted file mode 100644 index febdbfe39b..0000000000 --- a/packages/core/src/core/nonStrictFakeContentGenerator.ts +++ /dev/null @@ -1,111 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - GenerateContentResponse, - type CountTokensResponse, - type GenerateContentParameters, - type CountTokensParameters, - EmbedContentResponse, - type EmbedContentParameters, -} from '@google/genai'; -import { promises } from 'node:fs'; -import type { ContentGenerator } from './contentGenerator.js'; -import type { UserTierId, GeminiUserTier } from '../code_assist/types.js'; -import { safeJsonStringify } from '../utils/safeJsonStringify.js'; -import type { LlmRole } from '../telemetry/types.js'; -import type { FakeResponse } from './fakeContentGenerator.js'; - -/** - * A ContentGenerator that responds with canned responses, but unlike FakeContentGenerator, - * it is "non-strict": it will find and use the first available response that matches - * the requested method, rather than strictly following the input order. - * - * This is useful for testing asynchronous or non-deterministic background tasks - * (like token calibration or background snapshots) that might fire out-of-order. - */ -export class NonStrictFakeContentGenerator implements ContentGenerator { - userTier?: UserTierId; - userTierName?: string; - paidTier?: GeminiUserTier; - - constructor(private readonly responses: FakeResponse[]) {} - - static async fromFile( - filePath: string, - ): Promise { - const fileContent = await promises.readFile(filePath, 'utf-8'); - const responses = fileContent - .split('\n') - .filter((line) => line.trim() !== '') - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - .map((line) => JSON.parse(line) as FakeResponse); - return new NonStrictFakeContentGenerator(responses); - } - - private getNextResponse< - M extends FakeResponse['method'], - R = Extract['response'], - >(method: M, request: unknown): R { - const index = this.responses.findIndex((r) => r.method === method); - if (index === -1) { - throw new Error( - `No more mock responses for ${method}, got request:\n` + - safeJsonStringify(request), - ); - } - const response = this.responses.splice(index, 1)[0]; - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - return response.response as R; - } - - async generateContent( - request: GenerateContentParameters, - _userPromptId: string, - // eslint-disable-next-line @typescript-eslint/no-unused-vars - role: LlmRole, - ): Promise { - // eslint-disable-next-line @typescript-eslint/no-unsafe-return - return Object.setPrototypeOf( - this.getNextResponse('generateContent', request), - GenerateContentResponse.prototype, - ); - } - - async generateContentStream( - request: GenerateContentParameters, - _userPromptId: string, - // eslint-disable-next-line @typescript-eslint/no-unused-vars - role: LlmRole, - ): Promise> { - const responses = this.getNextResponse('generateContentStream', request); - async function* stream() { - for (const response of responses) { - yield Object.setPrototypeOf( - response, - GenerateContentResponse.prototype, - ); - } - } - return stream(); - } - - async countTokens( - request: CountTokensParameters, - ): Promise { - return this.getNextResponse('countTokens', request); - } - - async embedContent( - request: EmbedContentParameters, - ): Promise { - // eslint-disable-next-line @typescript-eslint/no-unsafe-return - return Object.setPrototypeOf( - this.getNextResponse('embedContent', request), - EmbedContentResponse.prototype, - ); - } -}