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

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
+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();
});