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'; import { vi } from 'vitest';
// Mock core modules // Mock core modules
vi.mock('./ui/hooks/atCommandProcessor.js');
vi.mock('@google/gemini-cli-core', async (importOriginal) => { vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const original = const original =
await importOriginal<typeof import('@google/gemini-cli-core')>(); await importOriginal<typeof import('@google/gemini-cli-core')>();
@ -41,7 +42,7 @@ describe('runNonInteractive', () => {
sendMessageStream: vi.Mock; sendMessageStream: vi.Mock;
}; };
beforeEach(() => { beforeEach(async () => {
mockCoreExecuteToolCall = vi.mocked(executeToolCall); mockCoreExecuteToolCall = vi.mocked(executeToolCall);
mockShutdownTelemetry = vi.mocked(shutdownTelemetry); mockShutdownTelemetry = vi.mocked(shutdownTelemetry);
@ -72,6 +73,14 @@ describe('runNonInteractive', () => {
getContentGeneratorConfig: vi.fn().mockReturnValue({}), getContentGeneratorConfig: vi.fn().mockReturnValue({}),
getDebugMode: vi.fn().mockReturnValue(false), getDebugMode: vi.fn().mockReturnValue(false),
} as unknown as Config; } as unknown as Config;
const { handleAtCommand } = await import(
'./ui/hooks/atCommandProcessor.js'
);
vi.mocked(handleAtCommand).mockImplementation(async ({ query }) => ({
processedQuery: [{ text: query }],
shouldProceed: true,
}));
}); });
afterEach(() => { afterEach(() => {
@ -163,14 +172,16 @@ describe('runNonInteractive', () => {
mockCoreExecuteToolCall.mockResolvedValue({ mockCoreExecuteToolCall.mockResolvedValue({
error: new Error('Execution failed'), error: new Error('Execution failed'),
errorType: ToolErrorType.EXECUTION_FAILED, errorType: ToolErrorType.EXECUTION_FAILED,
responseParts: { responseParts: [
functionResponse: { {
name: 'errorTool', functionResponse: {
response: { name: 'errorTool',
output: 'Error: Execution failed', response: {
output: 'Error: Execution failed',
},
}, },
}, },
}, ],
resultDisplay: 'Execution failed', resultDisplay: 'Execution failed',
}); });
const finalResponse: ServerGeminiStreamEvent[] = [ 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.', '\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 { Content, Part, FunctionCall } from '@google/genai';
import { ConsolePatcher } from './ui/utils/ConsolePatcher.js'; import { ConsolePatcher } from './ui/utils/ConsolePatcher.js';
import { handleAtCommand } from './ui/hooks/atCommandProcessor.js';
export async function runNonInteractive( export async function runNonInteractive(
config: Config, config: Config,
@ -40,9 +41,27 @@ export async function runNonInteractive(
const geminiClient = config.getGeminiClient(); const geminiClient = config.getGeminiClient();
const abortController = new AbortController(); 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[] = [ let currentMessages: Content[] = [
{ role: 'user', parts: [{ text: input }] }, { role: 'user', parts: processedQuery as Part[] },
]; ];
let turnCount = 0; let turnCount = 0;
while (true) { while (true) {
turnCount++; turnCount++;