diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index ed26f634a0..2a47036486 100644 --- a/packages/core/src/agents/local-executor.ts +++ b/packages/core/src/agents/local-executor.ts @@ -323,7 +323,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), @@ -810,6 +810,7 @@ export class LocalAgentExecutor { private async tryCompressChat( chat: GeminiChat, prompt_id: string, + abortSignal?: AbortSignal, ): Promise { const model = this.definition.modelConfig.model ?? DEFAULT_GEMINI_MODEL; @@ -820,6 +821,7 @@ export class LocalAgentExecutor { model, this.context.config, this.hasFailedCompressionAttempt, + abortSignal, ); if ( diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index f357a0decb..443a663219 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -608,7 +608,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 }; @@ -1158,6 +1158,7 @@ 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` @@ -1171,6 +1172,7 @@ export class GeminiClient { model, this.config, this.hasFailedCompressionAttempt, + abortSignal, ); if (