diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 2bb2836595..cfe8bdf34b 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -303,6 +303,17 @@ describe('Gemini Client (client.ts)', () => { }); }); + describe('setHistory', () => { + it('should update telemetry token count when history is set', () => { + const history: Content[] = [ + { role: 'user', parts: [{ text: 'some message' }] }, + ]; + client.setHistory(history); + + expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalled(); + }); + }); + describe('resumeChat', () => { it('should update telemetry token count when a chat is resumed', async () => { const history: Content[] = [ diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index fdf5e22a4d..fad1b8cbab 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -246,6 +246,7 @@ export class GeminiClient { setHistory(history: Content[]) { this.getChat().setHistory(history); + this.updateTelemetryTokenCount(); this.forceFullIdeContext = true; } diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 2bf22d509c..741e369f58 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -219,6 +219,37 @@ describe('GeminiChat', () => { }); }); + describe('setHistory', () => { + it('should recalculate lastPromptTokenCount when history is updated', () => { + const initialHistory: Content[] = [ + { role: 'user', parts: [{ text: 'Hello' }] }, + ]; + const chatWithHistory = new GeminiChat( + mockConfig, + '', + [], + initialHistory, + ); + const initialCount = chatWithHistory.getLastPromptTokenCount(); + + const newHistory: Content[] = [ + { + role: 'user', + parts: [ + { + text: 'This is a much longer history item that should result in more tokens than just hello.', + }, + ], + }, + ]; + chatWithHistory.setHistory(newHistory); + + expect(chatWithHistory.getLastPromptTokenCount()).toBeGreaterThan( + initialCount, + ); + }); + }); + describe('sendMessageStream', () => { it('should succeed if a tool call is followed by an empty part', async () => { // 1. Mock a stream that contains a tool call, then an invalid (empty) part. diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 4a99cf7b7b..bd7182fd03 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -676,6 +676,9 @@ export class GeminiChat { setHistory(history: Content[]): void { this.history = history; + this.lastPromptTokenCount = estimateTokenCountSync( + this.history.flatMap((c) => c.parts || []), + ); } stripThoughtsFromHistory(): void {