At Command Race Condition Bugfix For Non-Interactive Mode (#6676)

This commit is contained in:
Victor May 2025-08-21 14:47:40 -04:00 committed by GitHub
parent 1e5ead6960
commit 720eb81890
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 82 additions and 8 deletions

View File

@ -18,6 +18,7 @@ import { runNonInteractive } from './nonInteractiveCli.js';
import { vi } from 'vitest';
// Mock core modules
vi.mock('./ui/hooks/atCommandProcessor.js');
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const original =
await importOriginal<typeof import('@google/gemini-cli-core')>();
@ -41,7 +42,7 @@ describe('runNonInteractive', () => {
sendMessageStream: vi.Mock;
};
beforeEach(() => {
beforeEach(async () => {
mockCoreExecuteToolCall = vi.mocked(executeToolCall);
mockShutdownTelemetry = vi.mocked(shutdownTelemetry);
@ -72,6 +73,14 @@ describe('runNonInteractive', () => {
getContentGeneratorConfig: vi.fn().mockReturnValue({}),
getDebugMode: vi.fn().mockReturnValue(false),
} as unknown as Config;
const { handleAtCommand } = await import(
'./ui/hooks/atCommandProcessor.js'
);
vi.mocked(handleAtCommand).mockImplementation(async ({ query }) => ({
processedQuery: [{ text: query }],
shouldProceed: true,
}));
});
afterEach(() => {
@ -163,14 +172,16 @@ describe('runNonInteractive', () => {
mockCoreExecuteToolCall.mockResolvedValue({
error: new Error('Execution failed'),
errorType: ToolErrorType.EXECUTION_FAILED,
responseParts: {
functionResponse: {
name: 'errorTool',
response: {
output: 'Error: Execution failed',
responseParts: [
{
functionResponse: {
name: 'errorTool',
response: {
output: 'Error: Execution failed',
},
},
},
},
],
resultDisplay: 'Execution failed',
});
const finalResponse: ServerGeminiStreamEvent[] = [
@ -273,4 +284,48 @@ describe('runNonInteractive', () => {
'\n Reached max session turns for this session. Increase the number of turns by specifying maxSessionTurns in settings.json.',
);
});
it('should preprocess @include commands before sending to the model', async () => {
// 1. Mock the imported atCommandProcessor
const { handleAtCommand } = await import(
'./ui/hooks/atCommandProcessor.js'
);
const mockHandleAtCommand = vi.mocked(handleAtCommand);
// 2. Define the raw input and the expected processed output
const rawInput = 'Summarize @file.txt';
const processedParts: Part[] = [
{ text: 'Summarize @file.txt' },
{ text: '\n--- Content from referenced files ---\n' },
{ text: 'This is the content of the file.' },
{ text: '\n--- End of content ---' },
];
// 3. Setup the mock to return the processed parts
mockHandleAtCommand.mockResolvedValue({
processedQuery: processedParts,
shouldProceed: true,
});
// Mock a simple stream response from the Gemini client
const events: ServerGeminiStreamEvent[] = [
{ type: GeminiEventType.Content, value: 'Summary complete.' },
];
mockGeminiClient.sendMessageStream.mockReturnValue(
createStreamFromEvents(events),
);
// 4. Run the non-interactive mode with the raw input
await runNonInteractive(mockConfig, rawInput, 'prompt-id-7');
// 5. Assert that sendMessageStream was called with the PROCESSED parts, not the raw input
expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledWith(
processedParts,
expect.any(AbortSignal),
'prompt-id-7',
);
// 6. Assert the final output is correct
expect(processStdoutSpy).toHaveBeenCalledWith('Summary complete.');
});
});

View File

@ -16,6 +16,7 @@ import {
import { Content, Part, FunctionCall } from '@google/genai';
import { ConsolePatcher } from './ui/utils/ConsolePatcher.js';
import { handleAtCommand } from './ui/hooks/atCommandProcessor.js';
export async function runNonInteractive(
config: Config,
@ -40,9 +41,27 @@ export async function runNonInteractive(
const geminiClient = config.getGeminiClient();
const abortController = new AbortController();
const { processedQuery, shouldProceed } = await handleAtCommand({
query: input,
config,
addItem: (_item, _timestamp) => 0,
onDebugMessage: () => {},
messageId: Date.now(),
signal: abortController.signal,
});
if (!shouldProceed || !processedQuery) {
// An error occurred during @include processing (e.g., file not found).
// The error message is already logged by handleAtCommand.
console.error('Exiting due to an error processing the @ command.');
process.exit(1);
}
let currentMessages: Content[] = [
{ role: 'user', parts: [{ text: input }] },
{ role: 'user', parts: processedQuery as Part[] },
];
let turnCount = 0;
while (true) {
turnCount++;