bug(ux): update context percentage when /clear command is run (#4162)

Co-authored-by: matt korwel <[email protected]>

4 files changed, 160 insertions, 7 deletions

diff --git a/packages/cli/src/ui/commands/clearCommand.test.ts b/packages/cli/src/ui/commands/clearCommand.test.ts
index 8019dd68..10d54f4b 100644
--- a/packages/cli/src/ui/commands/clearCommand.test.ts
+++ b/packages/cli/src/ui/commands/clearCommand.test.ts
@@ -8,7 +8,19 @@ import { vi, describe, it, expect, beforeEach, Mock } from 'vitest';
 import { clearCommand } from './clearCommand.js';
 import { type CommandContext } from './types.js';
 import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
-import { GeminiClient } from '@google/gemini-cli-core';
+
+// Mock the telemetry service
+vi.mock('@google/gemini-cli-core', async () => {
+  const actual = await vi.importActual('@google/gemini-cli-core');
+  return {
+    ...actual,
+    uiTelemetryService: {
+      resetLastPromptTokenCount: vi.fn(),
+    },
+  };
+});
+
+import { GeminiClient, uiTelemetryService } from '@google/gemini-cli-core';
 
 describe('clearCommand', () => {
   let mockContext: CommandContext;
@@ -16,6 +28,7 @@ describe('clearCommand', () => {
 
   beforeEach(() => {
     mockResetChat = vi.fn().mockResolvedValue(undefined);
+    vi.clearAllMocks();
 
     mockContext = createMockCommandContext({
       services: {
@@ -29,7 +42,7 @@ describe('clearCommand', () => {
     });
   });
 
-  it('should set debug message, reset chat, and clear UI when config is available', async () => {
+  it('should set debug message, reset chat, reset telemetry, and clear UI when config is available', async () => {
     if (!clearCommand.action) {
       throw new Error('clearCommand must have an action.');
     }
@@ -42,18 +55,24 @@ describe('clearCommand', () => {
     expect(mockContext.ui.setDebugMessage).toHaveBeenCalledTimes(1);
 
     expect(mockResetChat).toHaveBeenCalledTimes(1);
-
+    expect(uiTelemetryService.resetLastPromptTokenCount).toHaveBeenCalledTimes(
+      1,
+    );
     expect(mockContext.ui.clear).toHaveBeenCalledTimes(1);
 
     // Check the order of operations.
     const setDebugMessageOrder = (mockContext.ui.setDebugMessage as Mock).mock
       .invocationCallOrder[0];
     const resetChatOrder = mockResetChat.mock.invocationCallOrder[0];
+    const resetTelemetryOrder = (
+      uiTelemetryService.resetLastPromptTokenCount as Mock
+    ).mock.invocationCallOrder[0];
     const clearOrder = (mockContext.ui.clear as Mock).mock
       .invocationCallOrder[0];
 
     expect(setDebugMessageOrder).toBeLessThan(resetChatOrder);
-    expect(resetChatOrder).toBeLessThan(clearOrder);
+    expect(resetChatOrder).toBeLessThan(resetTelemetryOrder);
+    expect(resetTelemetryOrder).toBeLessThan(clearOrder);
   });
 
   it('should not attempt to reset chat if config service is not available', async () => {
@@ -70,9 +89,12 @@ describe('clearCommand', () => {
     await clearCommand.action(nullConfigContext, '');
 
     expect(nullConfigContext.ui.setDebugMessage).toHaveBeenCalledWith(
-      'Clearing terminal and resetting chat.',
+      'Clearing terminal.',
     );
     expect(mockResetChat).not.toHaveBeenCalled();
+    expect(uiTelemetryService.resetLastPromptTokenCount).toHaveBeenCalledTimes(
+      1,
+    );
     expect(nullConfigContext.ui.clear).toHaveBeenCalledTimes(1);
   });
 });
diff --git a/packages/cli/src/ui/commands/clearCommand.ts b/packages/cli/src/ui/commands/clearCommand.ts
index e5473b5b..1c409359 100644
--- a/packages/cli/src/ui/commands/clearCommand.ts
+++ b/packages/cli/src/ui/commands/clearCommand.ts
@@ -4,14 +4,25 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { uiTelemetryService } from '@google/gemini-cli-core';
 import { SlashCommand } from './types.js';
 
 export const clearCommand: SlashCommand = {
   name: 'clear',
   description: 'clear the screen and conversation history',
   action: async (context, _args) => {
-    context.ui.setDebugMessage('Clearing terminal and resetting chat.');
-    await context.services.config?.getGeminiClient()?.resetChat();
+    const geminiClient = context.services.config?.getGeminiClient();
+
+    if (geminiClient) {
+      context.ui.setDebugMessage('Clearing terminal and resetting chat.');
+      // If resetChat fails, the exception will propagate and halt the command,
+      // which is the correct behavior to signal a failure to the user.
+      await geminiClient.resetChat();
+    } else {
+      context.ui.setDebugMessage('Clearing terminal.');
+    }
+
+    uiTelemetryService.resetLastPromptTokenCount();
     context.ui.clear();
   },
 };
diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts
index 34a2fe22..38ba7a91 100644
--- a/packages/core/src/telemetry/uiTelemetry.test.ts
+++ b/packages/core/src/telemetry/uiTelemetry.test.ts
@@ -508,4 +508,116 @@ describe('UiTelemetryService', () => {
       expect(tools.byName['tool_B'].count).toBe(1);
     });
   });
+
+  describe('resetLastPromptTokenCount', () => {
+    it('should reset the last prompt token count to 0', () => {
+      // First, set up some initial token count
+      const event = {
+        'event.name': EVENT_API_RESPONSE,
+        model: 'gemini-2.5-pro',
+        duration_ms: 500,
+        input_token_count: 100,
+        output_token_count: 200,
+        total_token_count: 300,
+        cached_content_token_count: 50,
+        thoughts_token_count: 20,
+        tool_token_count: 30,
+      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
+
+      service.addEvent(event);
+      expect(service.getLastPromptTokenCount()).toBe(100);
+
+      // Now reset the token count
+      service.resetLastPromptTokenCount();
+      expect(service.getLastPromptTokenCount()).toBe(0);
+    });
+
+    it('should emit an update event when resetLastPromptTokenCount is called', () => {
+      const spy = vi.fn();
+      service.on('update', spy);
+
+      // Set up initial token count
+      const event = {
+        'event.name': EVENT_API_RESPONSE,
+        model: 'gemini-2.5-pro',
+        duration_ms: 500,
+        input_token_count: 100,
+        output_token_count: 200,
+        total_token_count: 300,
+        cached_content_token_count: 50,
+        thoughts_token_count: 20,
+        tool_token_count: 30,
+      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
+
+      service.addEvent(event);
+      spy.mockClear(); // Clear the spy to focus on the reset call
+
+      service.resetLastPromptTokenCount();
+
+      expect(spy).toHaveBeenCalledOnce();
+      const { metrics, lastPromptTokenCount } = spy.mock.calls[0][0];
+      expect(metrics).toBeDefined();
+      expect(lastPromptTokenCount).toBe(0);
+    });
+
+    it('should not affect other metrics when resetLastPromptTokenCount is called', () => {
+      // Set up initial state with some metrics
+      const event = {
+        'event.name': EVENT_API_RESPONSE,
+        model: 'gemini-2.5-pro',
+        duration_ms: 500,
+        input_token_count: 100,
+        output_token_count: 200,
+        total_token_count: 300,
+        cached_content_token_count: 50,
+        thoughts_token_count: 20,
+        tool_token_count: 30,
+      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
+
+      service.addEvent(event);
+
+      const metricsBefore = service.getMetrics();
+
+      service.resetLastPromptTokenCount();
+
+      const metricsAfter = service.getMetrics();
+
+      // Metrics should be unchanged
+      expect(metricsAfter).toEqual(metricsBefore);
+
+      // Only the last prompt token count should be reset
+      expect(service.getLastPromptTokenCount()).toBe(0);
+    });
+
+    it('should work correctly when called multiple times', () => {
+      const spy = vi.fn();
+      service.on('update', spy);
+
+      // Set up initial token count
+      const event = {
+        'event.name': EVENT_API_RESPONSE,
+        model: 'gemini-2.5-pro',
+        duration_ms: 500,
+        input_token_count: 100,
+        output_token_count: 200,
+        total_token_count: 300,
+        cached_content_token_count: 50,
+        thoughts_token_count: 20,
+        tool_token_count: 30,
+      } as ApiResponseEvent & { 'event.name': typeof EVENT_API_RESPONSE };
+
+      service.addEvent(event);
+      expect(service.getLastPromptTokenCount()).toBe(100);
+
+      // Reset once
+      service.resetLastPromptTokenCount();
+      expect(service.getLastPromptTokenCount()).toBe(0);
+
+      // Reset again - should still be 0 and still emit event
+      spy.mockClear();
+      service.resetLastPromptTokenCount();
+      expect(service.getLastPromptTokenCount()).toBe(0);
+      expect(spy).toHaveBeenCalledOnce();
+    });
+  });
 });
diff --git a/packages/core/src/telemetry/uiTelemetry.ts b/packages/core/src/telemetry/uiTelemetry.ts
index 71409696..2713ac65 100644
--- a/packages/core/src/telemetry/uiTelemetry.ts
+++ b/packages/core/src/telemetry/uiTelemetry.ts
@@ -133,6 +133,14 @@ export class UiTelemetryService extends EventEmitter {
     return this.#lastPromptTokenCount;
   }
 
+  resetLastPromptTokenCount(): void {
+    this.#lastPromptTokenCount = 0;
+    this.emit('update', {
+      metrics: this.#metrics,
+      lastPromptTokenCount: this.#lastPromptTokenCount,
+    });
+  }
+
   private getOrCreateModelMetrics(modelName: string): ModelMetrics {
     if (!this.#metrics.models[modelName]) {
       this.#metrics.models[modelName] = createInitialModelMetrics();