diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 513424ad32a..f37cd58fc47 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -239,7 +239,7 @@ export class LocalAgentExecutor { ): Promise { const promptId = `${this.agentId}#${turnCounter}`; - await this.tryCompressChat(chat, promptId); + await this.tryCompressChat(chat, promptId, combinedSignal); const { functionCalls } = await promptIdContext.run(promptId, async () => this.callModel(chat, currentMessage, combinedSignal, promptId), @@ -668,6 +668,7 @@ export class LocalAgentExecutor { private async tryCompressChat( chat: GeminiChat, prompt_id: string, + abortSignal: AbortSignal, ): Promise { const model = this.definition.modelConfig.model ?? DEFAULT_GEMINI_MODEL; @@ -678,6 +679,7 @@ export class LocalAgentExecutor { model, this.runtimeContext, this.hasFailedCompressionAttempt, + abortSignal, ); if ( diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index c94dd5c04d0..9f1ce6c2ac0 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -574,7 +574,7 @@ export class GeminiClient { // Check for context window overflow const modelForLimitCheck = this._getActiveModelForCurrentTurn(); - const compressed = await this.tryCompressChat(prompt_id, false); + const compressed = await this.tryCompressChat(prompt_id, false, signal); if (compressed.compressionStatus === CompressionStatus.COMPRESSED) { yield { type: GeminiEventType.ChatCompressed, value: compressed }; @@ -1049,12 +1049,17 @@ export class GeminiClient { async tryCompressChat( prompt_id: string, force: boolean = false, + abortSignal?: AbortSignal, ): Promise { // If the model is 'auto', we will use a placeholder model to check. // Compression occurs before we choose a model, so calling `count_tokens` // before the model is chosen would result in an error. const model = this._getActiveModelForCurrentTurn(); + // Use the provided signal, or create a no-op signal as a fallback for + // callers that don't have one (e.g. the /compress command). + const signal = abortSignal ?? new AbortController().signal; + const { newHistory, info } = await this.compressionService.compress( this.getChat(), prompt_id, @@ -1062,6 +1067,7 @@ export class GeminiClient { model, this.config, this.hasFailedCompressionAttempt, + signal, ); if ( diff --git a/packages/core/src/services/chatCompressionService.test.ts b/packages/core/src/services/chatCompressionService.test.ts index 4ddd38e25cb..4bb42b9ff71 100644 --- a/packages/core/src/services/chatCompressionService.test.ts +++ b/packages/core/src/services/chatCompressionService.test.ts @@ -138,6 +138,7 @@ describe('ChatCompressionService', () => { let testTempDir: string; const mockModel = 'gemini-2.5-pro'; const mockPromptId = 'test-prompt-id'; + const mockAbortSignal = new AbortController().signal; beforeEach(() => { testTempDir = fs.mkdtempSync( @@ -211,6 +212,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.NOOP); expect(result.newHistory).toBeNull(); @@ -227,6 +229,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); // It should now attempt compression even if previously failed (logic removed) // But since history is small, it will be NOOP due to threshold @@ -253,6 +256,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.NOOP); expect(result.newHistory).toBeNull(); @@ -276,6 +280,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); @@ -317,6 +322,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); @@ -344,6 +350,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); const firstCall = vi.mocked(mockConfig.getBaseLlmClient().generateContent) @@ -371,6 +378,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); @@ -408,6 +416,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe( @@ -448,6 +457,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe( @@ -504,6 +514,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); @@ -570,6 +581,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); // Verify it compressed @@ -636,6 +648,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.newHistory).not.toBeNull(); @@ -703,6 +716,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); @@ -760,6 +774,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); @@ -817,6 +832,7 @@ describe('ChatCompressionService', () => { mockModel, mockConfig, false, + mockAbortSignal, ); expect(result.info.compressionStatus).toBe(CompressionStatus.COMPRESSED); diff --git a/packages/core/src/services/chatCompressionService.ts b/packages/core/src/services/chatCompressionService.ts index 5303a1a82a8..90549c01b0c 100644 --- a/packages/core/src/services/chatCompressionService.ts +++ b/packages/core/src/services/chatCompressionService.ts @@ -236,7 +236,7 @@ export class ChatCompressionService { model: string, config: Config, hasFailedCompressionAttempt: boolean, - abortSignal?: AbortSignal, + abortSignal: AbortSignal, ): Promise<{ newHistory: Content[] | null; info: ChatCompressionInfo }> { const curatedHistory = chat.getHistory(true); @@ -365,8 +365,7 @@ export class ChatCompressionService { ], systemInstruction: { text: getCompressionPrompt(config) }, promptId, - // TODO(joshualitt): wire up a sensible abort signal, - abortSignal: abortSignal ?? new AbortController().signal, + abortSignal, role: LlmRole.UTILITY_COMPRESSOR, }); const summary = getResponseText(summaryResponse) ?? ''; @@ -395,7 +394,7 @@ export class ChatCompressionService { systemInstruction: { text: getCompressionPrompt(config) }, promptId: `${promptId}-verify`, role: LlmRole.UTILITY_COMPRESSOR, - abortSignal: abortSignal ?? new AbortController().signal, + abortSignal, }); const finalSummary = (