From cc25ce89201610689e1fcd1644946b7dcfb6664d Mon Sep 17 00:00:00 2001 From: Spencer Date: Fri, 30 Jan 2026 22:48:03 +0000 Subject: [PATCH] feat: implement real-time quota visibility in CLI footer - Intercept rate-limit headers from both public (GoogleGenAI) and internal (CodeAssist) APIs. - Add QuotaDisplay UI component with color-coded remaining percentage. - Sync CodeAssistServer constructor with main (userTierName support). - Fix circular dependencies in contentGenerator.ts. - Improve test stability in clipboardUtils.test.ts by isolating environment state. --- packages/cli/src/ui/components/Footer.tsx | 2 +- .../cli/src/ui/utils/clipboardUtils.test.ts | 4 ++++ packages/cli/src/ui/utils/clipboardUtils.ts | 8 ++++++++ .../core/src/code_assist/codeAssist.test.ts | 7 ++++--- packages/core/src/code_assist/codeAssist.ts | 2 +- packages/core/src/code_assist/server.test.ts | 1 + .../core/src/core/contentGenerator.test.ts | 20 +++++++++++-------- 7 files changed, 31 insertions(+), 13 deletions(-) diff --git a/packages/cli/src/ui/components/Footer.tsx b/packages/cli/src/ui/components/Footer.tsx index c675df4dab..fa1934da12 100644 --- a/packages/cli/src/ui/components/Footer.tsx +++ b/packages/cli/src/ui/components/Footer.tsx @@ -197,4 +197,4 @@ export const Footer: React.FC = () => { )} ); -}; \ No newline at end of file +}; diff --git a/packages/cli/src/ui/utils/clipboardUtils.test.ts b/packages/cli/src/ui/utils/clipboardUtils.test.ts index 76eb0bcac3..ce710df88b 100644 --- a/packages/cli/src/ui/utils/clipboardUtils.test.ts +++ b/packages/cli/src/ui/utils/clipboardUtils.test.ts @@ -55,6 +55,7 @@ import { cleanupOldClipboardImages, splitEscapedPaths, parsePastedPaths, + resetClipboardToolForTesting, } from './clipboardUtils.js'; // Define the type for the module to use in tests @@ -73,6 +74,7 @@ describe('clipboardUtils', () => { process.env = { ...originalEnv }; // Reset modules to clear internal state (linuxClipboardTool variable) + resetClipboardToolForTesting(); vi.resetModules(); // Dynamically import the module to get a fresh instance for each test clipboardUtils = await import('./clipboardUtils.js'); @@ -305,6 +307,8 @@ describe('clipboardUtils', () => { }); it('should return null if tool is not yet detected', async () => { + setPlatform('linux'); + delete process.env['XDG_SESSION_TYPE']; // Don't prime the tool const result = await clipboardUtils.saveClipboardImage(mockTargetDir); expect(result).toBe(null); diff --git a/packages/cli/src/ui/utils/clipboardUtils.ts b/packages/cli/src/ui/utils/clipboardUtils.ts index 99ead45736..f1948a474d 100644 --- a/packages/cli/src/ui/utils/clipboardUtils.ts +++ b/packages/cli/src/ui/utils/clipboardUtils.ts @@ -35,6 +35,14 @@ const PATH_PREFIX_PATTERN = /^([/~.]|[a-zA-Z]:|\\\\)/; // Track which tool works on Linux to avoid redundant checks/failures let linuxClipboardTool: 'wl-paste' | 'xclip' | null = null; +/** + * Resets the cached Linux clipboard tool. + * Only intended for use in unit tests. + */ +export function resetClipboardToolForTesting() { + linuxClipboardTool = null; +} + // Helper to check the user's display server and whether they have a compatible clipboard tool installed function getUserLinuxClipboardTool(): typeof linuxClipboardTool { if (linuxClipboardTool !== null) { diff --git a/packages/core/src/code_assist/codeAssist.test.ts b/packages/core/src/code_assist/codeAssist.test.ts index b13c4be718..63aefdc3f0 100644 --- a/packages/core/src/code_assist/codeAssist.test.ts +++ b/packages/core/src/code_assist/codeAssist.test.ts @@ -43,6 +43,7 @@ describe('codeAssist', () => { const mockUserData = { projectId: 'test-project', userTier: UserTierId.FREE, + userTierName: 'free-tier-name', }; it('should create a server for LOGIN_WITH_GOOGLE', async () => { @@ -70,7 +71,7 @@ describe('codeAssist', () => { httpOptions, 'session-123', 'free-tier', - undefined, + 'free-tier-name', expect.any(Function), ); expect(generator).toBeInstanceOf(MockedCodeAssistServer); @@ -100,7 +101,7 @@ describe('codeAssist', () => { httpOptions, undefined, // No session ID 'free-tier', - undefined, + 'free-tier-name', expect.any(Function), ); expect(generator).toBeInstanceOf(MockedCodeAssistServer); @@ -173,4 +174,4 @@ describe('codeAssist', () => { expect(server).toBeUndefined(); }); }); -}); \ No newline at end of file +}); diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts index 52cf521618..79f1550e91 100644 --- a/packages/core/src/code_assist/codeAssist.ts +++ b/packages/core/src/code_assist/codeAssist.ts @@ -53,4 +53,4 @@ export function getCodeAssistServer( return undefined; } return server; -} \ No newline at end of file +} diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts index 6976b77ed5..7b07711efe 100644 --- a/packages/core/src/code_assist/server.test.ts +++ b/packages/core/src/code_assist/server.test.ts @@ -623,6 +623,7 @@ describe('CodeAssistServer', () => { {}, 'test-session', UserTierId.FREE, + undefined, onQuotaUpdate, ); diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 9ab65dd8e0..92eca0690a 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -368,11 +368,12 @@ describe('createContentGenerator', () => { expect(GoogleGenAI).toHaveBeenCalledWith({ apiKey: 'test-api-key', vertexai: undefined, - httpOptions: { + httpOptions: expect.objectContaining({ headers: expect.objectContaining({ 'User-Agent': expect.any(String), }), - }, + fetch: expect.any(Function), + }), apiVersion: 'v1', }); }); @@ -401,11 +402,12 @@ describe('createContentGenerator', () => { expect(GoogleGenAI).toHaveBeenCalledWith({ apiKey: 'test-api-key', vertexai: undefined, - httpOptions: { + httpOptions: expect.objectContaining({ headers: expect.objectContaining({ 'User-Agent': expect.any(String), }), - }, + fetch: expect.any(Function), + }), }); expect(GoogleGenAI).toHaveBeenCalledWith( @@ -440,11 +442,12 @@ describe('createContentGenerator', () => { expect(GoogleGenAI).toHaveBeenCalledWith({ apiKey: 'test-api-key', vertexai: undefined, - httpOptions: { + httpOptions: expect.objectContaining({ headers: expect.objectContaining({ 'User-Agent': expect.any(String), }), - }, + fetch: expect.any(Function), + }), }); expect(GoogleGenAI).toHaveBeenCalledWith( @@ -480,11 +483,12 @@ describe('createContentGenerator', () => { expect(GoogleGenAI).toHaveBeenCalledWith({ apiKey: 'test-api-key', vertexai: true, - httpOptions: { + httpOptions: expect.objectContaining({ headers: expect.objectContaining({ 'User-Agent': expect.any(String), }), - }, + fetch: expect.any(Function), + }), apiVersion: 'v1alpha', }); });