Agentic UI: Warn in the composer at 90% AI credit usage and complete the exhausted state

apps/ui/src/data/queries/use-agent-run.test.tsx

@@ -15,6 +15,12 @@ vi.mock( '@/data/core', async ( importOriginal ) => {
 	};
 } );
 
+const { outOfCreditsState } = vi.hoisted( () => ( { outOfCreditsState: { value: false } } ) );
+
+vi.mock( '@/hooks/use-is-out-of-ai-credits', () => ( {
+	useIsOutOfAiCredits: () => outOfCreditsState.value,
+} ) );
+
 const useConnectorMock = vi.mocked( useConnector );
 
 function createQueryClient() {
@@ -71,6 +77,7 @@ describe( 'useAgentRun queued handoff', () => {
 			} ),
 		};
 		useConnectorMock.mockReturnValue( connector as Connector );
+		outOfCreditsState.value = false;
 	} );
 
 	afterEach( () => {
@@ -141,6 +148,46 @@ describe( 'useAgentRun queued handoff', () => {
 		).toBe( true );
 	} );
 
+	it( 'holds a queued prompt instead of dispatching it once the credits are spent', async () => {
+		const queryClient = createQueryClient();
+		queryClient.setQueryData< LoadedAiSession >(
+			[ ...SESSIONS_QUERY_KEY, 'session-1' ],
+			createLoadedSession()
+		);
+
+		renderWithAgentRun( queryClient );
+		await waitFor( () => expect( connector.onAgentEvent ).toHaveBeenCalled() );
+
+		act( () => {
+			agentListener( {
+				sessionId: 'session-1',
+				runId: 'run-old',
+				event: { type: 'run.started', timestamp: '2026-06-24T12:00:00.000Z' },
+			} );
+		} );
+		await waitFor( () => expect( screen.getByTestId( 'phase' ) ).toHaveTextContent( 'active' ) );
+
+		// Queued while the run was still paid for; the balance empties mid-run.
+		fireEvent.click( screen.getByRole( 'button', { name: 'Queue' } ) );
+		outOfCreditsState.value = true;
+
+		act( () => {
+			agentListener( {
+				sessionId: 'session-1',
+				runId: 'run-old',
+				event: {
+					type: 'run.exited',
+					timestamp: '2026-06-24T12:00:01.000Z',
+					status: 'success',
+					code: 0,
+				},
+			} );
+		} );
+
+		await waitFor( () => expect( screen.getByTestId( 'phase' ) ).toHaveTextContent( 'idle' ) );
+		expect( connector.continueSession ).not.toHaveBeenCalled();
+	} );
+
 	it( 'still invalidates when a run ends without a queued follow-up', async () => {
 		const queryClient = createQueryClient();
 		const invalidateSpy = vi.spyOn( queryClient, 'invalidateQueries' );

apps/ui/src/data/queries/use-agent-run.tsx

@@ -16,6 +16,7 @@ import {
 import { useConnector } from '@/data/core';
 import { ASSISTANT_QUOTA_QUERY_KEY } from '@/data/queries/use-assistant-quota';
 import { SESSIONS_QUERY_KEY } from '@/data/queries/use-sessions';
+import { useIsOutOfAiCredits } from '@/hooks/use-is-out-of-ai-credits';
 import type {
 	AgentEvent,
 	AgentRunEvent,
@@ -725,6 +726,8 @@ export function useAgentRun( sessionId: string | undefined ): LiveAgentEvents {
 		queuedPrompts,
 	} = state;
 
+	const isOutOfCredits = useIsOutOfAiCredits();
+
 	// Re-entry guard for the queue auto-dispatch effect. The effect's deps
 	// re-fire on every queue/phase change; without this guard a second render
 	// between the async start-call and `send_start` could kick off a duplicate
@@ -735,7 +738,10 @@ export function useAgentRun( sessionId: string | undefined ): LiveAgentEvents {
 	// success, shift; on failure, drop the whole queue so a broken backend
 	// doesn't cascade errors.
 	useEffect( () => {
-		if ( ! sessionId || phase !== 'idle' || queuedPrompts.length === 0 ) {
+		// Holding rather than clearing: a balance spent mid-run leaves the queued
+		// prompt visible and it dispatches on its own once a top-up refreshes the
+		// quota, instead of being silently dropped or sent to a certain rejection.
+		if ( ! sessionId || phase !== 'idle' || queuedPrompts.length === 0 || isOutOfCredits ) {
 			return;
 		}
 		if ( dispatchingQueuedRef.current ) {
@@ -757,7 +763,7 @@ export function useAgentRun( sessionId: string | undefined ): LiveAgentEvents {
 				dispatchingQueuedRef.current = false;
 			}
 		} )();
-	}, [ dispatchSession, phase, queuedPrompts, sessionId, startRun ] );
+	}, [ dispatchSession, isOutOfCredits, phase, queuedPrompts, sessionId, startRun ] );
 
 	const sendMessage = useCallback(
 		async ( prompt: string, options: SendMessageOptions = {} ) => {

apps/ui/src/hooks/use-ai-credits-meter.test.ts

@@ -0,0 +1,53 @@
+import { renderHook } from '@testing-library/react';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { useStudioAssistantQuota } from '@/data/queries/use-assistant-quota';
+import { useAiCreditsMeter } from './use-ai-credits-meter';
+
+vi.mock( '@/data/queries/use-assistant-quota', () => ( {
+	useStudioAssistantQuota: vi.fn(),
+} ) );
+
+const useQuotaMock = vi.mocked( useStudioAssistantQuota );
+
+function mockQuota( data: unknown ) {
+	useQuotaMock.mockReturnValue( { data } as never );
+}
+
+describe( 'useAiCreditsMeter', () => {
+	beforeEach( () => vi.clearAllMocks() );
+
+	it( 'measures the free allowance against the monthly cap', () => {
+		mockQuota( { costUsage: 0, costCap: 1000000, allowanceRemaining: 100000 } );
+
+		expect( renderHook( () => useAiCreditsMeter() ).result.current ).toMatchObject( {
+			usedCredits: 900000,
+			totalCredits: 1000000,
+			remainingCredits: 100000,
+			fraction: 0.9,
+		} );
+	} );
+
+	it( 'is null while the quota is unknown', () => {
+		mockQuota( undefined );
+
+		expect( renderHook( () => useAiCreditsMeter() ).result.current ).toBeNull();
+	} );
+
+	it( 'is null for accounts the server reports no credit pools for', () => {
+		mockQuota( { costUsage: 25, costCap: 100 } );
+
+		expect( renderHook( () => useAiCreditsMeter() ).result.current ).toBeNull();
+	} );
+
+	it( 'is null for an account held by another access gate', () => {
+		mockQuota( {
+			costUsage: 0,
+			costCap: 1000000,
+			allowanceRemaining: 100000,
+			studioCodeAiHasAccess: false,
+			studioCodeAiAccess: 'blocked',
+		} );
+
+		expect( renderHook( () => useAiCreditsMeter() ).result.current ).toBeNull();
+	} );
+} );

apps/ui/src/hooks/use-ai-credits-meter.ts

@@ -0,0 +1,15 @@
+import {
+	getAiCreditsMeter,
+	getStudioCodeAiAccessState,
+	type AiCreditsMeter,
+} from '@studio/common/lib/studio-assistant-quota';
+import { useStudioAssistantQuota } from '@/data/queries/use-assistant-quota';
+
+/** Null means there is no figure to show, not an empty balance — see `useIsOutOfAiCredits` for that. */
+export function useAiCreditsMeter(): AiCreditsMeter | null {
+	const { data: quota } = useStudioAssistantQuota();
+	if ( ! quota || getStudioCodeAiAccessState( quota ) !== 'available' ) {
+		return null;
+	}
+	return getAiCreditsMeter( quota );
+}

apps/ui/src/ui-classic/components/session-view/composer/ai-credits-warning-strip.test.tsx

@@ -0,0 +1,74 @@
+import '@testing-library/jest-dom/vitest';
+import { render, screen } from '@testing-library/react';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { useStudioAssistantQuota } from '@/data/queries/use-assistant-quota';
+import { AiCreditsWarningStrip } from './ai-credits-warning-strip';
+
+vi.mock( '@/data/queries/use-assistant-quota', () => ( {
+	useStudioAssistantQuota: vi.fn(),
+} ) );
+
+vi.mock( '@/data/queries/use-user-locale', () => ( {
+	useUserLocale: () => 'en',
+} ) );
+
+vi.mock( '@/components/add-ai-credits-button', () => ( {
+	AddAiCreditsButton: () => <button type="button">Add AI credits</button>,
+} ) );
+
+const useQuotaMock = vi.mocked( useStudioAssistantQuota );
+
+function mockQuota( data: unknown ) {
+	useQuotaMock.mockReturnValue( { data } as never );
+}
+
+describe( 'AiCreditsWarningStrip', () => {
+	beforeEach( () => vi.clearAllMocks() );
+
+	it( 'interrupts the composer once the balance enters its last tenth', () => {
+		mockQuota( { costUsage: 0, costCap: 1000000, allowanceRemaining: 100000 } );
+
+		render( <AiCreditsWarningStrip /> );
+
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( 'At 90% usage' );
+		expect( screen.getByRole( 'button', { name: 'Add AI credits' } ) ).toBeInTheDocument();
+	} );
+
+	it( 'stays out of the way below the threshold', () => {
+		mockQuota( { costUsage: 0, costCap: 1000000, allowanceRemaining: 500000 } );
+
+		const { container } = render( <AiCreditsWarningStrip /> );
+
+		expect( container ).toBeEmptyDOMElement();
+	} );
+
+	it( 'yields to the lockout once the balance is spent', () => {
+		mockQuota( {
+			costUsage: 0,
+			costCap: 1000000,
+			allowanceRemaining: 0,
+			purchasedRemaining: 0,
+		} );
+
+		const { container } = render( <AiCreditsWarningStrip /> );
+
+		expect( container ).toBeEmptyDOMElement();
+	} );
+
+	it( 'renders nothing when there is no denominator to measure against', () => {
+		mockQuota( { costUsage: 25, costCap: 100 } );
+
+		const { container } = render( <AiCreditsWarningStrip /> );
+
+		expect( container ).toBeEmptyDOMElement();
+	} );
+
+	it( 'still warns when the meter rounds to full but credits remain', () => {
+		// The meter rounds up to full while the server balance is still above zero.
+		mockQuota( { costUsage: 0, costCap: 1000000, allowanceRemaining: 1 } );
+
+		render( <AiCreditsWarningStrip /> );
+
+		expect( screen.getByRole( 'status' ) ).toHaveTextContent( 'At 100% usage' );
+	} );
+} );

apps/ui/src/ui-classic/components/session-view/composer/ai-credits-warning-strip.tsx

@@ -0,0 +1,38 @@
+import { getAiCreditsMeterIntent } from '@studio/common/lib/studio-assistant-quota';
+import { __, sprintf } from '@wordpress/i18n';
+import { AddAiCreditsButton } from '@/components/add-ai-credits-button';
+import { useUserLocale } from '@/data/queries/use-user-locale';
+import { useAiCreditsMeter } from '@/hooks/use-ai-credits-meter';
+import { useIsOutOfAiCredits } from '@/hooks/use-is-out-of-ai-credits';
+import styles from './style.module.css';
+
+/** Warns in the composer from 90% usage. The lockout owns 100%; the two never render together. */
+export function AiCreditsWarningStrip() {
+	const meter = useAiCreditsMeter();
+	const isOutOfCredits = useIsOutOfAiCredits();
+	const locale = useUserLocale();
+	// 'exhausted' too: the meter can round up to full while credits remain.
+	const intent = meter ? getAiCreditsMeterIntent( meter.fraction ) : 'ok';
+	if ( ! meter || isOutOfCredits || ( intent !== 'critical' && intent !== 'exhausted' ) ) {
+		return null;
+	}
+
+	// Formatted, not concatenated: the percent sign moves and changes by locale.
+	const percentage = new Intl.NumberFormat( locale, {
+		style: 'percent',
+		maximumFractionDigits: 0,
+	} ).format( meter.fraction );
+
+	return (
+		<section className={ styles.aiCreditsWarningStrip } role="status">
+			<span>
+				{ sprintf(
+					/* translators: %s: share of the AI credit balance used, formatted as a percentage (e.g. 90%). */
+					__( 'At %s usage' ),
+					percentage
+				) }
+			</span>
+			<AddAiCreditsButton className={ styles.aiCreditsWarningStripButton } />
+		</section>
+	);
+}

apps/ui/src/ui-classic/components/session-view/composer/index.test.tsx

@@ -207,6 +207,30 @@ describe( 'Composer menu', () => {
 		expect( await screen.findByText( 'Stop' ) ).toBeInTheDocument();
 	} );
 
+	it( 'blocks sending and queueing while submission is unavailable', async () => {
+		const onSend = vi.fn< ( prompt: string ) => Promise< void > >();
+		renderComposer( { busy: true, canSubmit: false, onSend } );
+
+		const textarea = screen.getByRole( 'combobox' );
+		fireEvent.change( textarea, { target: { value: 'sneak one past the lockout' } } );
+
+		// The Enter path reaches send() directly, bypassing the button's state.
+		fireEvent.keyDown( textarea, { key: 'Enter' } );
+		fireEvent.click( screen.getByRole( 'button', { name: 'Queue' } ) );
+
+		await waitFor( () => expect( screen.getByRole( 'button', { name: 'Queue' } ) ).toBeDisabled() );
+		expect( onSend ).not.toHaveBeenCalled();
+	} );
+
+	it( 'keeps Stop working while submission is unavailable', () => {
+		const onInterrupt = vi.fn< () => Promise< void > >();
+		renderComposer( { busy: true, canSubmit: false, onInterrupt } );
+
+		fireEvent.click( screen.getByRole( 'button', { name: 'Stop' } ) );
+
+		expect( onInterrupt ).toHaveBeenCalledTimes( 1 );
+	} );
+
 	it( 'uses a queue-focused placeholder while busy', () => {
 		renderComposer( { busy: true } );
 

apps/ui/src/ui-classic/components/session-view/composer/index.tsx

@@ -66,6 +66,7 @@ import {
 	SESSIONS_QUERY_KEY,
 } from '@/data/queries/use-sessions';
 import { AiCreditsControl } from './ai-credits-control';
+import { AiCreditsWarningStrip } from './ai-credits-warning-strip';
 import { clearComposerDraft, getComposerDraft, saveComposerDraft } from './draft-store';
 import { FamilySwitchConfirmDialog } from './family-switch-confirm-dialog';
 import styles from './style.module.css';
@@ -238,6 +239,9 @@ export function ComposerSkeleton() {
 
 interface ComposerProps {
 	busy: boolean;
+	// Blocks sending and queueing while leaving the rest of the composer alone,
+	// so a run already in flight keeps its Stop control.
+	canSubmit?: boolean;
 	isInterrupting?: boolean;
 	error: string | null;
 	model: AiModelId;
@@ -328,6 +332,7 @@ function resizeComposerTextarea(
 const ComposerContent = forwardRef< ComposerHandle, ComposerProps >( function ComposerContent(
 	{
 		busy,
+		canSubmit = true,
 		isInterrupting = false,
 		error,
 		model,
@@ -497,6 +502,11 @@ const ComposerContent = forwardRef< ComposerHandle, ComposerProps >( function Co
 	);
 
 	const send = useCallback( async () => {
+		// Guarded here as well as on the button: Enter reaches this directly, and
+		// while busy a send becomes a queued prompt that would dispatch later.
+		if ( ! canSubmit ) {
+			return;
+		}
 		const trimmed = value.trim();
 		// Allow sending attachments on their own; fall back to a minimal prompt so
 		// the backend (which requires a non-empty message) still has one.
@@ -532,6 +542,7 @@ const ComposerContent = forwardRef< ComposerHandle, ComposerProps >( function Co
 			restoreAttachments( sentAttachments );
 		}
 	}, [
+		canSubmit,
 		value,
 		attachments,
 		suggestionBaseline,
@@ -758,7 +769,7 @@ const ComposerContent = forwardRef< ComposerHandle, ComposerProps >( function Co
 		}
 	}, [ connector, onSwitchSession, ownerSiteId, pendingFamilyChange, queryClient ] );
 
-	const canSend = value.trim().length > 0 || attachments.length > 0;
+	const canSend = canSubmit && ( value.trim().length > 0 || attachments.length > 0 );
 	const placeholderOptions = busy
 		? [
 				__( 'Queue the next message while I work…' ),
@@ -811,6 +822,7 @@ const ComposerContent = forwardRef< ComposerHandle, ComposerProps >( function Co
 					onDragLeave={ dragHandlers.onDragLeave }
 					onDrop={ dragHandlers.onDrop }
 				>
+					<AiCreditsWarningStrip />
 					<div
 						className={ styles.resizeHandle }
 						role="separator"