diff --git a/packages/core/src/agents/local-executor.ts b/packages/core/src/agents/local-executor.ts index 513424ad32a..43f366af40b 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..9e34f5239a5 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,6 +1049,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` @@ -1062,6 +1063,7 @@ export class GeminiClient { model, this.config, this.hasFailedCompressionAttempt, + abortSignal, ); if (