AI credits: Warn at 80% usage in the sidebar and above the Classic composer

apps/studio/src/components/ai-credits-threshold-notice.tsx

@@ -0,0 +1,91 @@
+import {
+	formatAiCreditsThresholdDescription,
+	formatAiCreditsUsageTitle,
+	getAiCreditsMeter,
+	getAiCreditsMeterIntent,
+	getStudioCodeAiAccessState,
+	resolveAiCreditsThresholdNotice,
+} from '@studio/common/lib/studio-assistant-quota';
+import { __ } from '@wordpress/i18n';
+import { close } from '@wordpress/icons';
+import { Icon } from '@wordpress/ui';
+import { useEffect, type ReactNode } from 'react';
+import { AddAiCreditsButton } from 'src/components/add-ai-credits-button';
+import { useAppDispatch, useI18nLocale, useRootSelector } from 'src/stores';
+import { selectDismissedAiCreditsIntent, setDismissedAiCreditsIntent } from 'src/stores/ui-slice';
+import { useGetStudioAssistantQuota } from 'src/stores/wpcom-api';
+
+// Classic has no composer strip and no lockout banner of its own, so this one
+// slot announces both warning steps. The agentic UI splits them across the
+// sidebar and the composer instead.
+const CLASSIC_NOTICE_INTENTS = [ 'warning', 'critical' ] as const;
+
+function NoticeCard( {
+	title,
+	description,
+	action,
+	onDismiss,
+}: {
+	title: string;
+	description?: string;
+	action?: ReactNode;
+	onDismiss: () => void;
+} ) {
+	return (
+		<div className="border-frame-border bg-frame-surface relative mb-2 flex flex-col items-start gap-1 rounded-lg border p-3 text-left">
+			<span className="text-frame-text pe-6 text-sm font-semibold">{ title }</span>
+			{ description ? (
+				<span className="text-frame-text-secondary text-xs">{ description }</span>
+			) : null }
+			{ action }
+			<button
+				type="button"
+				aria-label={ __( 'Dismiss' ) }
+				onClick={ onDismiss }
+				className="text-frame-text-secondary hover:text-frame-text absolute end-2 top-2"
+			>
+				<Icon icon={ close } size={ 16 } />
+			</button>
+		</div>
+	);
+}
+
+/**
+ * Warns above the Classic composer as the AI credit balance runs down. Classic
+ * has no persistent-message surface, so the warning sits where the user is
+ * about to spend the credits.
+ */
+export function AiCreditsThresholdNotice() {
+	const dispatch = useAppDispatch();
+	const locale = useI18nLocale();
+	const dismissedIntent = useRootSelector( selectDismissedAiCreditsIntent );
+	const { data: quota } = useGetStudioAssistantQuota();
+	const meter =
+		quota && getStudioCodeAiAccessState( quota ) === 'available'
+			? getAiCreditsMeter( quota )
+			: null;
+	const intent = meter ? getAiCreditsMeterIntent( meter.fraction ) : null;
+	const notice = resolveAiCreditsThresholdNotice( intent, dismissedIntent, CLASSIC_NOTICE_INTENTS );
+
+	// Drop a dismissal the current usage has left behind, so the notice can
+	// fire again if the account returns to that threshold.
+	useEffect( () => {
+		if ( notice.dismissedIntent !== dismissedIntent ) {
+			dispatch( setDismissedAiCreditsIntent( notice.dismissedIntent ) );
+		}
+	}, [ dispatch, dismissedIntent, notice.dismissedIntent ] );
+
+	const thresholdIntent = intent === 'warning' || intent === 'critical' ? intent : null;
+	if ( ! thresholdIntent || ! notice.visible || ! meter ) {
+		return null;
+	}
+
+	return (
+		<NoticeCard
+			title={ formatAiCreditsUsageTitle( meter.fraction, locale ) }
+			description={ formatAiCreditsThresholdDescription() }
+			action={ <AddAiCreditsButton className="mt-2" /> }
+			onDismiss={ () => dispatch( setDismissedAiCreditsIntent( thresholdIntent ) ) }
+		/>
+	);
+}

apps/studio/src/components/studio-code-session/index.tsx

@@ -21,6 +21,7 @@ import {
 	type UIEvent,
 } from 'react';
 import { OutOfCreditsNotice } from 'src/components/ai-access-required-notice';
+import { AiCreditsThresholdNotice } from 'src/components/ai-credits-threshold-notice';
 import { ArrowIcon } from 'src/components/arrow-icon';
 import Button from 'src/components/button';
 import { IllustrationGrid } from 'src/components/illustration-grid';
@@ -422,6 +423,7 @@ function SessionContent( { selectedSite }: { selectedSite: SiteDetails } ) {
 				composer={
 					<div className={ styles.classicColumn }>
 						<QueuedPrompts prompts={ queuedPrompts } onRemove={ removeQueuedPrompt } />
+						<AiCreditsThresholdNotice />
 						{ isOutOfCredits ? (
 							<OutOfCreditsNotice />
 						) : (

apps/studio/src/components/studio-code-session/tests/access-requirements.test.tsx

@@ -31,6 +31,9 @@ vi.mock( 'src/hooks/use-auth', () => ( {
 
 vi.mock( 'src/stores', () => ( {
 	useI18nLocale: () => 'en',
+	useAppDispatch: () => vi.fn(),
+	// Neutral for every ui-slice selector the session tree reads.
+	useRootSelector: () => null,
 } ) );
 
 vi.mock( 'src/stores/wpcom-api', () => ( {

apps/studio/src/components/studio-code-session/tests/session-not-found.test.tsx

@@ -25,6 +25,9 @@ vi.mock( 'src/hooks/use-auth', () => ( {
 
 vi.mock( 'src/stores', () => ( {
 	useI18nLocale: () => 'en',
+	useAppDispatch: () => vi.fn(),
+	// Neutral for every ui-slice selector the session tree reads.
+	useRootSelector: () => null,
 } ) );
 
 vi.mock( 'src/stores/wpcom-api', () => ( {

apps/studio/src/components/tests/ai-credits-threshold-notice.test.tsx

@@ -0,0 +1,87 @@
+import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { AiCreditsThresholdNotice } from 'src/components/ai-credits-threshold-notice';
+import { useAppDispatch, useRootSelector } from 'src/stores';
+import { setDismissedAiCreditsIntent } from 'src/stores/ui-slice';
+import { useGetStudioAssistantQuota } from 'src/stores/wpcom-api';
+import type { AiCreditsMeterIntent } from '@studio/common/lib/studio-assistant-quota';
+
+vi.mock( 'src/stores', () => ( {
+	useAppDispatch: vi.fn(),
+	useRootSelector: vi.fn(),
+	useI18nLocale: () => 'en',
+} ) );
+vi.mock( 'src/stores/wpcom-api', () => ( { useGetStudioAssistantQuota: vi.fn() } ) );
+vi.mock( 'src/components/add-ai-credits-button', () => ( {
+	AddAiCreditsButton: () => <button type="button">Add AI credits</button>,
+} ) );
+
+const dispatch = vi.fn();
+
+// 1,000,000-credit allowance, nothing purchased.
+function mockUsage( allowanceRemaining: number ) {
+	vi.mocked( useGetStudioAssistantQuota, { partial: true } ).mockReturnValue( {
+		data: { costUsage: 0, costCap: 1000000, allowanceRemaining, purchasedRemaining: undefined },
+	} );
+}
+
+function mockState( {
+	dismissedIntent = null,
+}: { dismissedIntent?: AiCreditsMeterIntent | null } = {} ) {
+	vi.mocked( useRootSelector ).mockImplementation( ( selector ) =>
+		selector( { ui: { dismissedAiCreditsIntent: dismissedIntent } } as never )
+	);
+}
+
+describe( 'AiCreditsThresholdNotice', () => {
+	beforeEach( () => {
+		vi.clearAllMocks();
+		vi.mocked( useAppDispatch, { partial: true } ).mockReturnValue( dispatch );
+		mockState();
+	} );
+
+	it( 'warns at 80% usage, reporting the live figure', () => {
+		mockUsage( 170000 );
+		render( <AiCreditsThresholdNotice /> );
+
+		expect( screen.getByText( 'At 83% usage' ) ).toBeInTheDocument();
+		expect( screen.getByRole( 'button', { name: 'Add AI credits' } ) ).toBeInTheDocument();
+	} );
+
+	// Classic has no composer strip, so its one slot carries the 90% step too.
+	it( 'escalates at 90% usage', () => {
+		mockUsage( 100000 );
+		render( <AiCreditsThresholdNotice /> );
+
+		expect( screen.getByText( 'At 90% usage' ) ).toBeInTheDocument();
+	} );
+
+	it( 'shows nothing below 80%, or once the credits are spent', () => {
+		mockUsage( 500000 );
+		const { rerender } = render( <AiCreditsThresholdNotice /> );
+		expect( screen.queryByText( /usage$/ ) ).not.toBeInTheDocument();
+
+		mockUsage( 0 );
+		rerender( <AiCreditsThresholdNotice /> );
+		expect( screen.queryByText( /usage$/ ) ).not.toBeInTheDocument();
+	} );
+
+	it( 'records the dismissal against the threshold it was made at', async () => {
+		const user = userEvent.setup();
+		mockUsage( 200000 );
+		render( <AiCreditsThresholdNotice /> );
+
+		await user.click( screen.getByRole( 'button', { name: 'Dismiss' } ) );
+
+		expect( dispatch ).toHaveBeenCalledWith( setDismissedAiCreditsIntent( 'warning' ) );
+	} );
+
+	it( 'stays hidden at a threshold already dismissed', () => {
+		mockUsage( 200000 );
+		mockState( { dismissedIntent: 'warning' } );
+		render( <AiCreditsThresholdNotice /> );
+
+		expect( screen.queryByText( 'At 80% usage' ) ).not.toBeInTheDocument();
+	} );
+} );

apps/studio/src/stores/ui-slice.ts

@@ -1,14 +1,19 @@
-import { createSlice } from '@reduxjs/toolkit';
-import { RootState } from 'src/stores';
+import { createSlice, type PayloadAction } from '@reduxjs/toolkit';
+import type { AiCreditsMeterIntent } from '@studio/common/lib/studio-assistant-quota';
+import type { RootState } from 'src/stores';
 
 type UiState = {
 	isAddSiteModalOpen: boolean;
 	isWapuuWorldOpen: boolean;
+	// AI credits notice state, session-only on purpose: a threshold notice
+	// describes the balance right now, not a standing preference.
+	dismissedAiCreditsIntent: AiCreditsMeterIntent | null;
 };
 
 const initialState: UiState = {
 	isAddSiteModalOpen: false,
 	isWapuuWorldOpen: false,
+	dismissedAiCreditsIntent: null,
 };
 
 const uiSlice = createSlice( {
@@ -30,13 +35,26 @@ const uiSlice = createSlice( {
 		closeWapuuWorld: ( state ) => {
 			state.isWapuuWorldOpen = false;
 		},
+		setDismissedAiCreditsIntent: (
+			state,
+			action: PayloadAction< AiCreditsMeterIntent | null >
+		) => {
+			state.dismissedAiCreditsIntent = action.payload;
+		},
 	},
 } );
 
-export const { openAddSiteModal, closeAddSiteModal, openWapuuWorld, closeWapuuWorld } =
-	uiSlice.actions;
+export const {
+	openAddSiteModal,
+	closeAddSiteModal,
+	openWapuuWorld,
+	closeWapuuWorld,
+	setDismissedAiCreditsIntent,
+} = uiSlice.actions;
 
 export const selectIsAddSiteModalOpen = ( state: RootState ) => state.ui.isAddSiteModalOpen;
 export const selectIsWapuuWorldOpen = ( state: RootState ) => state.ui.isWapuuWorldOpen;
+export const selectDismissedAiCreditsIntent = ( state: RootState ) =>
+	state.ui.dismissedAiCreditsIntent;
 
 export default uiSlice.reducer;

apps/ui/src/components/app-message-cards/index.tsx

@@ -1,6 +1,7 @@
 import { __ } from '@wordpress/i18n';
 import { Button, Notice } from '@wordpress/ui';
 import { clsx } from 'clsx';
+import { AddAiCreditsButton } from '@/components/add-ai-credits-button';
 import toastStyles from '@/components/app-toasts/style.module.css';
 import { useActivePersistentMessages } from '@/data/queries/use-app-messages';
 import styles from './style.module.css';
@@ -25,7 +26,15 @@ export function AppMessageCards( { className }: { className?: string } ) {
 						{ message.description ? (
 							<Notice.Description>{ message.description }</Notice.Description>
 						) : null }
-						{ message.cta ? (
+						{ message.purchaseCta ? (
+							<Notice.Actions>
+								<AddAiCreditsButton
+									variant="solid"
+									tone="neutral"
+									className={ toastStyles.actionButton }
+								/>
+							</Notice.Actions>
+						) : message.cta ? (
 							<Notice.Actions>
 								<Button
 									size="small"

apps/ui/src/data/ai-credits-notice.test.ts

@@ -0,0 +1,25 @@
+import { renderHook, act } from '@testing-library/react';
+import { beforeEach, describe, expect, it } from 'vitest';
+import {
+	resetAiCreditsNoticeForTests,
+	setDismissedAiCreditsIntent,
+	useDismissedAiCreditsIntent,
+} from './ai-credits-notice';
+
+describe( 'ai credits notice dismissal', () => {
+	beforeEach( () => resetAiCreditsNoticeForTests() );
+
+	it( 'starts undismissed, so a fresh run always warns', () => {
+		expect( renderHook( () => useDismissedAiCreditsIntent() ).result.current ).toBeNull();
+	} );
+
+	it( 'notifies subscribers when the dismissal changes', () => {
+		const { result } = renderHook( () => useDismissedAiCreditsIntent() );
+
+		act( () => setDismissedAiCreditsIntent( 'warning' ) );
+		expect( result.current ).toBe( 'warning' );
+
+		act( () => setDismissedAiCreditsIntent( null ) );
+		expect( result.current ).toBeNull();
+	} );
+} );

apps/ui/src/data/ai-credits-notice.ts

@@ -0,0 +1,35 @@
+import { useSyncExternalStore } from 'react';
+import type { AiCreditsMeterIntent } from '@studio/common/lib/studio-assistant-quota';
+
+// Session-only, and deliberately not persisted: a threshold notice describes
+// the balance right now, so an account that reaches 80% again in a later run
+// has spent a new pool's worth of credits and deserves the warning again.
+let dismissedIntent: AiCreditsMeterIntent | null = null;
+const listeners = new Set< () => void >();
+
+export function setDismissedAiCreditsIntent( intent: AiCreditsMeterIntent | null ): void {
+	if ( dismissedIntent === intent ) {
+		return;
+	}
+	dismissedIntent = intent;
+	for ( const listener of listeners ) {
+		listener();
+	}
+}
+
+export function useDismissedAiCreditsIntent(): AiCreditsMeterIntent | null {
+	return useSyncExternalStore(
+		( listener ) => {
+			listeners.add( listener );
+			return () => {
+				listeners.delete( listener );
+			};
+		},
+		() => dismissedIntent,
+		() => null
+	);
+}
+
+export function resetAiCreditsNoticeForTests(): void {
+	dismissedIntent = null;
+}