MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-27 05:24:34 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix(ui): Track last prompt token count in processStream (#8650)

Browse Source
This commit is contained in:
Sandy Tao
2025-09-18 11:46:56 -07:00
committed by GitHub
parent bcfd50b45b
commit 2c754d71e3
8 changed files with 72 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/telemetry/uiTelemetry.test.ts
+11 -11
View File
@@ -142,7 +142,7 @@ describe('UiTelemetryService', () => {
expect(spy).toHaveBeenCalledOnce();
const { metrics, lastPromptTokenCount } = spy.mock.calls[0][0];
expect(metrics).toBeDefined();
expect(lastPromptTokenCount).toBe(10);
expect(lastPromptTokenCount).toBe(0);
});
describe('API Response Event Processing', () => {
@@ -177,7 +177,7 @@ describe('UiTelemetryService', () => {
tool: 3,
},
});
expect(service.getLastPromptTokenCount()).toBe(10);
expect(service.getLastPromptTokenCount()).toBe(0);
});
it('should aggregate multiple ApiResponseEvents for the same model', () => {
@@ -227,7 +227,7 @@ describe('UiTelemetryService', () => {
tool: 9,
},
});
expect(service.getLastPromptTokenCount()).toBe(15);
expect(service.getLastPromptTokenCount()).toBe(0);
});
it('should handle ApiResponseEvents for different models', () => {
@@ -266,7 +266,7 @@ describe('UiTelemetryService', () => {
expect(metrics.models['gemini-2.5-flash']).toBeDefined();
expect(metrics.models['gemini-2.5-pro'].api.totalRequests).toBe(1);
expect(metrics.models['gemini-2.5-flash'].api.totalRequests).toBe(1);
expect(service.getLastPromptTokenCount()).toBe(100);
expect(service.getLastPromptTokenCount()).toBe(0);
});
});
@@ -543,10 +543,10 @@ describe('UiTelemetryService', () => {
} as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
service.addEvent(event);
expect(service.getLastPromptTokenCount()).toBe(100);
expect(service.getLastPromptTokenCount()).toBe(0);
// Now reset the token count
service.resetLastPromptTokenCount();
service.setLastPromptTokenCount(0);
expect(service.getLastPromptTokenCount()).toBe(0);
});
@@ -570,7 +570,7 @@ describe('UiTelemetryService', () => {
service.addEvent(event);
spy.mockClear(); // Clear the spy to focus on the reset call
service.resetLastPromptTokenCount();
service.setLastPromptTokenCount(0);
expect(spy).toHaveBeenCalledOnce();
const { metrics, lastPromptTokenCount } = spy.mock.calls[0][0];
@@ -596,7 +596,7 @@ describe('UiTelemetryService', () => {
const metricsBefore = service.getMetrics();
service.resetLastPromptTokenCount();
service.setLastPromptTokenCount(0);
const metricsAfter = service.getMetrics();
@@ -625,15 +625,15 @@ describe('UiTelemetryService', () => {
} as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
service.addEvent(event);
expect(service.getLastPromptTokenCount()).toBe(100);
expect(service.getLastPromptTokenCount()).toBe(0);
// Reset once
service.resetLastPromptTokenCount();
service.setLastPromptTokenCount(0);
expect(service.getLastPromptTokenCount()).toBe(0);
// Reset again - should still be 0 and still emit event
spy.mockClear();
service.resetLastPromptTokenCount();
service.setLastPromptTokenCount(0);
expect(service.getLastPromptTokenCount()).toBe(0);
expect(spy).toHaveBeenCalledOnce();
});
packages/core/src/telemetry/uiTelemetry.ts
+2 -4
View File
@@ -144,8 +144,8 @@ export class UiTelemetryService extends EventEmitter {
return this.#lastPromptTokenCount;
}
resetLastPromptTokenCount(): void {
this.#lastPromptTokenCount = 0;
setLastPromptTokenCount(lastPromptTokenCount: number): void {
this.#lastPromptTokenCount = lastPromptTokenCount;
this.emit('update', {
metrics: this.#metrics,
lastPromptTokenCount: this.#lastPromptTokenCount,
@@ -171,8 +171,6 @@ export class UiTelemetryService extends EventEmitter {
modelMetrics.tokens.cached += event.cached_content_token_count;
modelMetrics.tokens.thoughts += event.thoughts_token_count;
modelMetrics.tokens.tool += event.tool_token_count;
this.#lastPromptTokenCount = event.input_token_count;
}
private processApiError(event: ApiErrorEvent) {