From 720eb81890c3d4b479accb851c77c4ee869d6024 Mon Sep 17 00:00:00 2001
From: Victor May <mayvic@google.com>
Date: Thu, 21 Aug 2025 14:47:40 -0400
Subject: [PATCH] At Command Race Condition Bugfix For Non-Interactive Mode
 (#6676)

---
 packages/cli/src/nonInteractiveCli.test.ts | 69 +++++++++++++++++++---
 packages/cli/src/nonInteractiveCli.ts      | 21 ++++++-
 2 files changed, 82 insertions(+), 8 deletions(-)

diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts
index d7f7ad70..8dc8775d 100644
--- a/packages/cli/src/nonInteractiveCli.test.ts
+++ b/packages/cli/src/nonInteractiveCli.test.ts
@@ -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.');
+  });
 });
diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts
index 6aec2754..36337c8f 100644
--- a/packages/cli/src/nonInteractiveCli.ts
+++ b/packages/cli/src/nonInteractiveCli.ts
@@ -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++;