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

View File

@@ -16,7 +16,7 @@ vi.mock('@google/gemini-cli-core', async () => {
return {
...actual,
uiTelemetryService: {
resetLastPromptTokenCount: vi.fn(),
setLastPromptTokenCount: vi.fn(),
},
};
});
@@ -57,9 +57,8 @@ describe('clearCommand', () => {
expect(mockContext.ui.setDebugMessage).toHaveBeenCalledTimes(1);
expect(mockResetChat).toHaveBeenCalledTimes(1);
expect(uiTelemetryService.resetLastPromptTokenCount).toHaveBeenCalledTimes(
1,
);
expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(0);
expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledTimes(1);
expect(mockContext.ui.clear).toHaveBeenCalledTimes(1);
// Check the order of operations.
@@ -67,7 +66,7 @@ describe('clearCommand', () => {
.invocationCallOrder[0];
const resetChatOrder = mockResetChat.mock.invocationCallOrder[0];
const resetTelemetryOrder = (
uiTelemetryService.resetLastPromptTokenCount as Mock
uiTelemetryService.setLastPromptTokenCount as Mock
).mock.invocationCallOrder[0];
const clearOrder = (mockContext.ui.clear as Mock).mock
.invocationCallOrder[0];
@@ -94,9 +93,8 @@ describe('clearCommand', () => {
'Clearing terminal.',
);
expect(mockResetChat).not.toHaveBeenCalled();
expect(uiTelemetryService.resetLastPromptTokenCount).toHaveBeenCalledTimes(
1,
);
expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(0);
expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledTimes(1);
expect(nullConfigContext.ui.clear).toHaveBeenCalledTimes(1);
});
});

View File

@@ -24,7 +24,7 @@ export const clearCommand: SlashCommand = {
context.ui.setDebugMessage('Clearing terminal.');
}
uiTelemetryService.resetLastPromptTokenCount();
uiTelemetryService.setLastPromptTokenCount(0);
context.ui.clear();
},
};