Add support for specifying maxSessionTurns via the settings configuration (#3507)

5 files changed, 80 insertions, 1 deletions

diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts
index b80b6dd0..b685f090 100644
--- a/packages/cli/src/config/config.ts
+++ b/packages/cli/src/config/config.ts
@@ -312,6 +312,7 @@ export async function loadCliConfig(
     bugCommand: settings.bugCommand,
     model: argv.model!,
     extensionContextFilePaths,
+    maxSessionTurns: settings.maxSessionTurns ?? -1,
     listExtensions: argv.listExtensions || false,
     activeExtensions: activeExtensions.map((e) => ({
       name: e.config.name,
diff --git a/packages/cli/src/config/settings.ts b/packages/cli/src/config/settings.ts
index 133701f5..2abe8cd8 100644
--- a/packages/cli/src/config/settings.ts
+++ b/packages/cli/src/config/settings.ts
@@ -80,6 +80,9 @@ export interface Settings {
   hideWindowTitle?: boolean;
   hideTips?: boolean;
 
+  // Setting for setting maximum number of user/model/tool turns in a session.
+  maxSessionTurns?: number;
+
   // Add other settings here.
 }
 
diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts
index 14352f53..6cbb630d 100644
--- a/packages/cli/src/nonInteractiveCli.test.ts
+++ b/packages/cli/src/nonInteractiveCli.test.ts
@@ -53,6 +53,7 @@ describe('runNonInteractive', () => {
       getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
       getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient),
       getContentGeneratorConfig: vi.fn().mockReturnValue({}),
+      getMaxSessionTurns: vi.fn().mockReturnValue(10),
       initialize: vi.fn(),
     } as unknown as Config;
 
@@ -294,4 +295,50 @@ describe('runNonInteractive', () => {
       'Unfortunately the tool does not exist.',
     );
   });
+
+  it('should exit when max session turns are exceeded', async () => {
+    const functionCall: FunctionCall = {
+      id: 'fcLoop',
+      name: 'loopTool',
+      args: {},
+    };
+    const toolResponsePart: Part = {
+      functionResponse: {
+        name: 'loopTool',
+        id: 'fcLoop',
+        response: { result: 'still looping' },
+      },
+    };
+
+    // Config with a max turn of 1
+    vi.mocked(mockConfig.getMaxSessionTurns).mockReturnValue(1);
+
+    const { executeToolCall: mockCoreExecuteToolCall } = await import(
+      '@google/gemini-cli-core'
+    );
+    vi.mocked(mockCoreExecuteToolCall).mockResolvedValue({
+      callId: 'fcLoop',
+      responseParts: [toolResponsePart],
+      resultDisplay: 'Still looping',
+      error: undefined,
+    });
+
+    const stream = (async function* () {
+      yield { functionCalls: [functionCall] } as GenerateContentResponse;
+    })();
+
+    mockChat.sendMessageStream.mockResolvedValue(stream);
+    const consoleErrorSpy = vi
+      .spyOn(console, 'error')
+      .mockImplementation(() => {});
+
+    await runNonInteractive(mockConfig, 'Trigger loop');
+
+    expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(1);
+    expect(consoleErrorSpy).toHaveBeenCalledWith(
+      `
+ Reached max session turns for this session. Increase the number of turns by specifying maxSessionTurns in settings.json.`,
+    );
+    expect(mockProcessExit).not.toHaveBeenCalled();
+  });
 });
diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts
index b8b8ac3f..2db28eba 100644
--- a/packages/cli/src/nonInteractiveCli.ts
+++ b/packages/cli/src/nonInteractiveCli.ts
@@ -63,9 +63,19 @@ export async function runNonInteractive(
   const chat = await geminiClient.getChat();
   const abortController = new AbortController();
   let currentMessages: Content[] = [{ role: 'user', parts: [{ text: input }] }];
-
+  let turnCount = 0;
   try {
     while (true) {
+      turnCount++;
+      if (
+        config.getMaxSessionTurns() > 0 &&
+        turnCount > config.getMaxSessionTurns()
+      ) {
+        console.error(
+          '\n Reached max session turns for this session. Increase the number of turns by specifying maxSessionTurns in settings.json.',
+        );
+        return;
+      }
       const functionCalls: FunctionCall[] = [];
 
       const responseStream = await chat.sendMessageStream(
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index b82b0cb2..a9326528 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -431,6 +431,20 @@ export const useGeminiStream = (
     [addItem, config],
   );
 
+  const handleMaxSessionTurnsEvent = useCallback(
+    () =>
+      addItem(
+        {
+          type: 'info',
+          text:
+            `The session has reached the maximum number of turns: ${config.getMaxSessionTurns()}. ` +
+            `Please update this limit in your setting.json file.`,
+        },
+        Date.now(),
+      ),
+    [addItem, config],
+  );
+
   const processGeminiStreamEvents = useCallback(
     async (
       stream: AsyncIterable<GeminiEvent>,
@@ -467,6 +481,9 @@ export const useGeminiStream = (
           case ServerGeminiEventType.ToolCallResponse:
             // do nothing
             break;
+          case ServerGeminiEventType.MaxSessionTurns:
+            handleMaxSessionTurnsEvent();
+            break;
           default: {
             // enforces exhaustive switch-case
             const unreachable: never = event;
@@ -485,6 +502,7 @@ export const useGeminiStream = (
       handleErrorEvent,
       scheduleToolCalls,
       handleChatCompressionEvent,
+      handleMaxSessionTurnsEvent,
     ],
   );