fix(core): keep token escalation warm across agent rounds

packages/core/src/agents/runtime/agent-core.ts

@@ -654,6 +654,7 @@ export class AgentCore {
     let turnCounter = 0;
     let finalText = '';
     let terminateMode: AgentTerminateMode | null = null;
+    let stickyMaxOutputTokens: number | undefined;
 
     while (true) {
       // Check abort before starting a new round — prevents unnecessary API
@@ -689,6 +690,9 @@ export class AgentCore {
           config: {
             abortSignal: roundAbortController.signal,
             tools: [{ functionDeclarations: toolsList }],
+            ...(stickyMaxOutputTokens !== undefined
+              ? { maxOutputTokens: stickyMaxOutputTokens }
+              : {}),
           },
         };
 
@@ -728,6 +732,9 @@ export class AgentCore {
           // retry does not inherit stale data (e.g. wasOutputTruncated) from a
           // previous attempt that may have hit MAX_TOKENS.
           if (streamEvent.type === 'retry') {
+            if (streamEvent.maxOutputTokensEscalated !== undefined) {
+              stickyMaxOutputTokens = streamEvent.maxOutputTokensEscalated;
+            }
             functionCalls.length = 0;
             roundText = '';
             roundThoughtText = '';

packages/core/src/agents/runtime/agent-headless.test.ts

@@ -1671,6 +1671,96 @@ describe('subagent.ts', () => {
           'rejected to prevent writing truncated content',
         );
       });
+
+      it('keeps automatic max token escalation warm for the next agent round', async () => {
+        const writeFileToolDef: FunctionDeclaration = {
+          name: WriteFileTool.Name,
+          description: 'Writes a file',
+          parameters: { type: Type.OBJECT, properties: {} },
+        };
+
+        const { config } = await createMockConfig({
+          getFunctionDeclarationsFiltered: vi
+            .fn()
+            .mockReturnValue([writeFileToolDef]),
+          getTool: vi.fn().mockImplementation((name: string) => {
+            if (name === WriteFileTool.Name) {
+              return new WriteFileTool(config);
+            }
+            return undefined;
+          }),
+        });
+
+        const toolConfig: ToolConfig = { tools: [WriteFileTool.Name] };
+        let callCount = 0;
+        mockSendMessageStream.mockImplementation(async () => {
+          callCount++;
+          if (callCount === 1) {
+            return (async function* () {
+              yield {
+                type: 'retry',
+                maxOutputTokensEscalated: 65_536,
+              };
+              yield {
+                type: 'chunk',
+                value: {
+                  functionCalls: [
+                    {
+                      id: 'call_write_complete',
+                      name: WriteFileTool.Name,
+                      args: {
+                        file_path: '/tmp/sticky-escalation.txt',
+                        content: 'hello',
+                      },
+                    },
+                  ],
+                },
+              };
+              yield {
+                type: 'chunk',
+                value: {
+                  candidates: [
+                    { finishReason: 'STOP', content: { parts: [] } },
+                  ],
+                },
+              };
+            })();
+          }
+
+          return (async function* () {
+            yield {
+              type: 'chunk',
+              value: {
+                candidates: [
+                  {
+                    finishReason: 'STOP',
+                    content: { parts: [{ text: 'done' }] },
+                  },
+                ],
+              },
+            };
+          })();
+        });
+
+        const scope = await AgentHeadless.create(
+          'test-agent',
+          config,
+          promptConfig,
+          defaultModelConfig,
+          defaultRunConfig,
+          toolConfig,
+        );
+
+        await scope.execute(new ContextState());
+
+        expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
+        expect(
+          mockSendMessageStream.mock.calls[0][1].config.maxOutputTokens,
+        ).toBeUndefined();
+        expect(
+          mockSendMessageStream.mock.calls[1][1].config.maxOutputTokens,
+        ).toBe(65_536);
+      });
     });
   });
 });

packages/core/src/core/geminiChat.test.ts

@@ -5908,6 +5908,37 @@ describe('GeminiChat', async () => {
       );
     });
 
+    it('should not re-escalate when the request already uses the escalated output limit', async () => {
+      vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(
+        makeStream([makeChunk([{ text: 'still partial' }], 'MAX_TOKENS')]),
+      );
+
+      const stream = await chat.sendMessageStream(
+        'gemini-3-pro',
+        {
+          message: 'continue the agent task',
+          config: { maxOutputTokens: 65_536 },
+        },
+        'prompt-sticky-escalation',
+      );
+
+      const events: StreamEvent[] = [];
+      for await (const event of stream) {
+        events.push(event);
+      }
+
+      expect(
+        events.some(
+          (event) =>
+            event.type === StreamEventType.RETRY &&
+            event.maxOutputTokensEscalated !== undefined,
+        ),
+      ).toBe(false);
+      expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(
+        1,
+      );
+    });
+
     it('should coalesce overlapping recovery continuation text', async () => {
       const streams = [
         makeStream([makeChunk([{ text: 'discarded initial' }], 'MAX_TOKENS')]),

packages/core/src/core/geminiChat.ts

@@ -208,6 +208,8 @@ export type StreamEvent =
        *  fresh restart (escalation). The UI should keep the accumulated text
        *  buffer so the continuation appends to it. */
       isContinuation?: boolean;
+      /** Set when the retry raised the automatic max output token limit. */
+      maxOutputTokensEscalated?: number;
     }
   | { type: StreamEventType.COMPRESSED; info: ChatCompressionInfo };
 
@@ -2202,19 +2204,29 @@ export class GeminiChat {
         // models.
         // Placed outside the retry loop so that any errors from the
         // escalated stream propagate directly (not caught by retry logic).
+        const requestedMaxOutputTokens = params.config?.maxOutputTokens;
+        const escalatedLimit = Math.max(
+          ESCALATED_MAX_TOKENS,
+          tokenLimit(model, 'output'),
+        );
+        const shouldEscalateMaxOutputTokens =
+          requestedMaxOutputTokens === undefined ||
+          requestedMaxOutputTokens < escalatedLimit;
+
         if (
           lastError === null &&
           lastFinishReason === FinishReason.MAX_TOKENS &&
           !maxTokensEscalated &&
-          !hasUserMaxTokensOverride
+          !hasUserMaxTokensOverride &&
+          shouldEscalateMaxOutputTokens
         ) {
           maxTokensEscalated = true;
-          const escalatedLimit = Math.max(
-            ESCALATED_MAX_TOKENS,
-            tokenLimit(model, 'output'),
-          );
+          const startingLimitLabel =
+            requestedMaxOutputTokens === undefined
+              ? 'capped default'
+              : `${requestedMaxOutputTokens} tokens`;
           debugLogger.info(
-            `Output truncated at capped default. Escalating to ${escalatedLimit} tokens.`,
+            `Output truncated at ${startingLimitLabel}. Escalating to ${escalatedLimit} tokens.`,
           );
           // Remove partial model response from history
           // (processStreamResponse already pushed it)
@@ -2225,7 +2237,10 @@ export class GeminiChat {
             self.history.pop();
           }
           // Signal UI to discard partial output
-          yield { type: StreamEventType.RETRY };
+          yield {
+            type: StreamEventType.RETRY,
+            maxOutputTokensEscalated: escalatedLimit,
+          };
           // Retry with escalated max_tokens
           const escalatedParams: SendMessageParameters = {
             ...params,