From a57b3a24deada9ff8ab7e39d89ce14ccdb0e0c2d Mon Sep 17 00:00:00 2001 From: Volodymyr Makukha Date: Tue, 16 Jun 2026 12:09:50 +0100 Subject: [PATCH] Add assistant cost usage limits to Settings --- .../user-settings/components/prompt-info.tsx | 59 +++++++++++++- .../components/tests/prompt-info.test.tsx | 81 +++++++++++++++++++ .../components/tests/user-settings.test.tsx | 3 +- apps/studio/src/stores/wpcom-api.ts | 25 ++++++ 4 files changed, 166 insertions(+), 2 deletions(-) create mode 100644 apps/studio/src/modules/user-settings/components/tests/prompt-info.test.tsx diff --git a/apps/studio/src/modules/user-settings/components/prompt-info.tsx b/apps/studio/src/modules/user-settings/components/prompt-info.tsx index 68b8371440..c524ed1824 100644 --- a/apps/studio/src/modules/user-settings/components/prompt-info.tsx +++ b/apps/studio/src/modules/user-settings/components/prompt-info.tsx @@ -1,7 +1,42 @@ +import { sprintf } from '@wordpress/i18n'; import { useI18n } from '@wordpress/react-i18n'; +import ProgressBar from 'src/components/progress-bar'; +import { useOffline } from 'src/hooks/use-offline'; +import { useI18nLocale } from 'src/stores'; +import { useGetStudioAssistantQuota } from 'src/stores/wpcom-api'; + +function formatPercentage( value: number, maxValue: number, locale: string ) { + const percentage = Math.max( 0, Math.min( 1, value / maxValue ) ); + + return new Intl.NumberFormat( locale, { + style: 'percent', + maximumFractionDigits: 2, + } ).format( percentage ); +} + +function formatResetDate( date: string, locale: string ) { + return new Intl.DateTimeFormat( locale, { + day: 'numeric', + month: 'long', + year: 'numeric', + } ).format( new Date( date ) ); +} export function PromptInfo() { const { __ } = useI18n(); + const locale = useI18nLocale(); + const isOffline = useOffline(); + const { + data: assistantQuota, + isError, + isLoading, + } = useGetStudioAssistantQuota( undefined, { + refetchOnMountOrArgChange: true, + } ); + const assistantQuotaWithCostCap = + assistantQuota && assistantQuota.costCap > 0 && ! isOffline && ! isError + ? assistantQuota + : undefined; return (
@@ -11,10 +46,32 @@ export function PromptInfo() {
- { __( 'Generous token limits while Studio Code is in beta.' ) } + { isOffline && __( "You're currently offline" ) } + { ! isOffline && isLoading && __( 'Loading Studio Code limits…' ) } + { assistantQuotaWithCostCap && + sprintf( + __( '%1$s of monthly limit used (resets on %2$s)' ), + formatPercentage( + assistantQuotaWithCostCap.costUsage, + assistantQuotaWithCostCap.costCap, + locale + ), + formatResetDate( assistantQuotaWithCostCap.costResetDate, locale ) + ) } + { ! isLoading && + ! isOffline && + ! assistantQuotaWithCostCap && + __( 'Studio Code limits are temporarily unavailable.' ) }
+ { ! isOffline && isLoading && } + { assistantQuotaWithCostCap && ( + + ) }
diff --git a/apps/studio/src/modules/user-settings/components/tests/prompt-info.test.tsx b/apps/studio/src/modules/user-settings/components/tests/prompt-info.test.tsx new file mode 100644 index 0000000000..5a670b28e4 --- /dev/null +++ b/apps/studio/src/modules/user-settings/components/tests/prompt-info.test.tsx @@ -0,0 +1,81 @@ +import { render, screen } from '@testing-library/react'; +import { vi } from 'vitest'; +import { useOffline } from 'src/hooks/use-offline'; +import { PromptInfo } from 'src/modules/user-settings/components/prompt-info'; +import { useGetStudioAssistantQuota } from 'src/stores/wpcom-api'; + +vi.mock( 'src/hooks/use-offline' ); + +vi.mock( 'src/stores', () => ( { + useI18nLocale: vi.fn( () => 'en-US' ), +} ) ); + +vi.mock( 'src/stores/wpcom-api', () => ( { + useGetStudioAssistantQuota: vi.fn(), +} ) ); + +describe( 'PromptInfo', () => { + beforeEach( () => { + vi.mocked( useOffline ).mockReturnValue( false ); + } ); + + it( 'shows Studio Code dollar usage and reset date', () => { + vi.mocked( useGetStudioAssistantQuota ).mockReturnValue( { + data: { + costUsage: 33392, + costCap: 20000000, + costResetDate: '2026-07-01T00:00:00+00:00', + }, + isError: false, + isLoading: false, + refetch: vi.fn(), + } as unknown as ReturnType< typeof useGetStudioAssistantQuota > ); + + render( ); + + expect( screen.getByText( 'Studio Code' ) ).toBeInTheDocument(); + expect( + screen.getByText( '0.17% of monthly limit used (resets on July 1, 2026)' ) + ).toBeInTheDocument(); + expect( screen.queryByText( /monthly prompts used/ ) ).not.toBeInTheDocument(); + expect( screen.getByRole( 'progressbar' ) ).toBeInTheDocument(); + } ); + + it( 'shows unavailable message when cost cap is missing', () => { + vi.mocked( useGetStudioAssistantQuota ).mockReturnValue( { + data: { + costUsage: 0, + costCap: 0, + costResetDate: '2026-07-01T00:00:00+00:00', + }, + isError: false, + isLoading: false, + refetch: vi.fn(), + } as unknown as ReturnType< typeof useGetStudioAssistantQuota > ); + + render( ); + + expect( + screen.getByText( 'Studio Code limits are temporarily unavailable.' ) + ).toBeInTheDocument(); + } ); + + it( 'caps over-limit usage at 100%', () => { + vi.mocked( useGetStudioAssistantQuota ).mockReturnValue( { + data: { + costUsage: 3403700000, + costCap: 20000000, + costResetDate: '2026-07-01T00:00:00+00:00', + }, + isError: false, + isLoading: false, + refetch: vi.fn(), + } as unknown as ReturnType< typeof useGetStudioAssistantQuota > ); + + render( ); + + expect( + screen.getByText( '100% of monthly limit used (resets on July 1, 2026)' ) + ).toBeInTheDocument(); + } ); +} ); diff --git a/apps/studio/src/modules/user-settings/components/tests/user-settings.test.tsx b/apps/studio/src/modules/user-settings/components/tests/user-settings.test.tsx index 7c70e56fbf..26ecf1f4c7 100644 --- a/apps/studio/src/modules/user-settings/components/tests/user-settings.test.tsx +++ b/apps/studio/src/modules/user-settings/components/tests/user-settings.test.tsx @@ -50,6 +50,7 @@ const mockIpcEvent = { describe( 'UserSettings', () => { beforeEach( () => { + vi.mocked( useOffline ).mockReturnValue( false ); // Triggers IPC listener to show modal vi.mocked( useIpcListener ).mockImplementationOnce( ( listener, callback ) => { if ( listener === 'user-settings' ) { @@ -143,7 +144,7 @@ describe( 'UserSettings', () => { expect( screen.getByText( 'Preview sites' ) ).toBeInTheDocument(); expect( screen.getByText( 'Studio Code' ) ).toBeInTheDocument(); expect( - screen.getByText( 'Generous token limits while Studio Code is in beta.' ) + screen.getByText( 'Studio Code limits are temporarily unavailable.' ) ).toBeInTheDocument(); expect( screen.queryByText( /monthly prompts used/ ) ).not.toBeInTheDocument(); } ); diff --git a/apps/studio/src/stores/wpcom-api.ts b/apps/studio/src/stores/wpcom-api.ts index 5ce34d9ac3..0bb8c52fe9 100644 --- a/apps/studio/src/stores/wpcom-api.ts +++ b/apps/studio/src/stores/wpcom-api.ts @@ -35,6 +35,18 @@ const snapshotStatusSchema = z isDeleted: data.is_deleted === '1', } ) ); +const studioAssistantQuotaSchema = z + .object( { + cost_usage: z.number(), + cost_cap: z.number(), + cost_reset_date: z.string(), + } ) + .transform( ( data ) => ( { + costUsage: data.cost_usage, + costCap: data.cost_cap, + costResetDate: data.cost_reset_date, + } ) ); + export type { Blueprint }; let wpcomClient: WPCOM | undefined; @@ -120,6 +132,15 @@ export const wpcomApi = createApi( { transformResponse: ( response: unknown ) => parseResponse( response, snapshotStatusSchema ), keepUnusedDataFor: 60 * 60, } ), + getStudioAssistantQuota: builder.query< z.infer< typeof studioAssistantQuotaSchema >, void >( { + query: () => ( { + path: '/studio-app/ai-assistant/quota', + apiNamespace: 'wpcom/v2', + } ), + transformResponse: ( response: unknown ) => + parseResponse( response, studioAssistantQuotaSchema ), + keepUnusedDataFor: 60 * 60, + } ), deleteAllSnapshots: builder.mutation< void, void >( { queryFn: async () => { await getIpcApi().deleteAllSnapshots(); @@ -219,6 +240,10 @@ export const useGetSnapshotStatus = withWpcomClientCheck( withOfflineCheck( wpcomApi.useGetSnapshotStatusQuery ) ); +export const useGetStudioAssistantQuota = withWpcomClientCheck( + withOfflineCheck( wpcomApi.useGetStudioAssistantQuotaQuery ) +); + export const useDeleteAllSnapshots = withWpcomClientCheckMutation( withOfflineCheckMutation( wpcomApi.useDeleteAllSnapshotsMutation ) );