Studio: Disable text input if the limit is reached UI and update to n…

…ew headers (#244)

* Studio: Fix quota headers based on new endpoint

* Studio: Disable input on reached quota

src/components/content-tab-assistant.tsx

@@ -9,6 +9,7 @@ import { useAssistant, Message as MessageType } from '../hooks/use-assistant';
 import { useAssistantApi } from '../hooks/use-assistant-api';
 import { useAuth } from '../hooks/use-auth';
 import { useOffline } from '../hooks/use-offline';
+import { usePromptUsage } from '../hooks/use-prompt-usage';
 import { cx } from '../lib/cx';
 import { getIpcApi } from '../lib/get-ipc-api';
 import { AIInput } from './ai-input';
@@ -233,6 +234,7 @@ const UnauthenticatedView = ( { onAuthenticate }: { onAuthenticate: () => void }
 
 export function ContentTabAssistant( { selectedSite }: ContentTabAssistantProps ) {
 	const { messages, addMessage, clearMessages } = useAssistant( selectedSite.name );
+	const { userCanSendMessage } = usePromptUsage();
 	const { fetchAssistant, isLoading: isAssistantThinking } = useAssistantApi();
 	const [ input, setInput ] = useState< string >( '' );
 	const endOfMessagesRef = useRef< HTMLDivElement >( null );
@@ -284,7 +286,7 @@ export function ContentTabAssistant( { selectedSite }: ContentTabAssistantProps
 		}
 	}, [ messages ] );
 
-	const disabled = isOffline || ! isAuthenticated;
+	const disabled = isOffline || ! isAuthenticated || ! userCanSendMessage;
 
 	return (
 		<div className="h-full flex flex-col bg-gray-50">

src/components/user-settings.tsx

@@ -3,7 +3,7 @@ import { sprintf } from '@wordpress/i18n';
 import { moreVertical, trash } from '@wordpress/icons';
 import { useI18n } from '@wordpress/react-i18n';
 import { useCallback, useState, useEffect } from 'react';
-import { WPCOM_PROFILE_URL } from '../constants';
+import { LIMIT_OF_PROMPTS_PER_USER, WPCOM_PROFILE_URL } from '../constants';
 import { useAuth } from '../hooks/use-auth';
 import { useDeleteSnapshot } from '../hooks/use-delete-snapshot';
 import { useFetchSnapshots } from '../hooks/use-fetch-snapshots';
@@ -142,7 +142,7 @@ const SnapshotInfo = ( {
 
 function PromptInfo() {
 	const { __ } = useI18n();
-	const { promptCount = 0, promptLimit = 10 } = usePromptUsage();
+	const { promptCount = 0, promptLimit = LIMIT_OF_PROMPTS_PER_USER } = usePromptUsage();
 	const assistantEnabled = getAppGlobals().assistantEnabled;
 	if ( ! assistantEnabled ) {
 		return null;

src/hooks/use-assistant-api.ts

@@ -48,7 +48,10 @@ export function useAssistantApi() {
 				setIsLoading( false );
 			}
 			const message = response?.choices?.[ 0 ]?.message?.content;
-			updatePromptUsage( headers );
+			updatePromptUsage( {
+				maxQuota: headers[ 'x-quota-max' ] || '',
+				remainingQuota: headers[ 'x-quota-remaining' ] || '',
+			} );
 			return { message };
 		},
 		[ client, updatePromptUsage ]

src/hooks/use-prompt-usage.tsx

@@ -8,14 +8,16 @@ type PromptUsage = {
 	promptLimit: number;
 	promptCount: number;
 	fetchPromptUsage: () => Promise< void >;
-	updatePromptUsage: ( headers: Record< string, string > ) => void;
+	updatePromptUsage: ( data: { maxQuota: string; remainingQuota: string } ) => void;
+	userCanSendMessage: boolean;
 };
 
 const initState = {
 	promptLimit: LIMIT_OF_PROMPTS_PER_USER,
 	promptCount: 0,
 	fetchPromptUsage: async () => undefined,
-	updatePromptUsage: ( _headers: Record< string, string > ) => undefined,
+	updatePromptUsage: ( _data: { maxQuota: string; remainingQuota: string } ) => undefined,
+	userCanSendMessage: true,
 };
 const promptUsageContext = createContext< PromptUsage >( initState );
 
@@ -37,9 +39,9 @@ export function PromptUsageProvider( { children }: PromptUsageProps ) {
 	const { client } = useAuth();
 
 	const updatePromptUsage = useCallback(
-		( headers: Record< string, string > ) => {
-			const limit = parseInt( headers[ 'x-ratelimit-limit' ] );
-			const remaining = parseInt( headers[ 'x-ratelimit-remaining' ] );
+		( data: { maxQuota: string; remainingQuota: string } ) => {
+			const limit = parseInt( data.maxQuota as string );
+			const remaining = parseInt( data.remainingQuota as string );
 			if ( isNaN( limit ) || isNaN( remaining ) ) {
 				return;
 			}
@@ -57,22 +59,15 @@ export function PromptUsageProvider( { children }: PromptUsageProps ) {
 			return;
 		}
 		try {
-			await client.req.get(
-				{
-					method: 'HEAD',
-					path: '/studio-app/ai-assistant/chat',
-					apiNamespace: 'wpcom/v2',
-				},
-				( error: Error, _data: unknown, headers: Record< string, string > ) => {
-					if ( error ) {
-						return;
-					}
-					if ( ! headers ) {
-						return;
-					}
-					updatePromptUsage( headers );
-				}
-			);
+			const response = await client.req.get( {
+				method: 'GET',
+				path: '/studio-app/ai-assistant/quota',
+				apiNamespace: 'wpcom/v2',
+			} );
+			updatePromptUsage( {
+				maxQuota: response.max_quota || '',
+				remainingQuota: response.remaining_quota || '',
+			} );
 		} catch ( error ) {
 			Sentry.captureException( error );
 			console.error( error );
@@ -92,6 +87,7 @@ export function PromptUsageProvider( { children }: PromptUsageProps ) {
 			promptLimit,
 			promptCount,
 			updatePromptUsage,
+			userCanSendMessage: promptCount < promptLimit,
 		};
 	}, [ fetchPromptUsage, promptLimit, promptCount, updatePromptUsage ] );