From 753c548f2983cb92859dbde6431328ccbf408c81 Mon Sep 17 00:00:00 2001
From: Adnan Haque <3737780+haqadn@users.noreply.github.com>
Date: Fri, 20 Dec 2024 17:38:43 -0500
Subject: [PATCH] Boost: Improve responsiveness during CCSS generation retry
 (#40675)

* Move C.CSS retry down the element chain

* Optimistic update critical CSS regen

* Only start generation if there is any provider in the list

* changelog

* Change the order of state update

* Fix showstopper behavior for cloud CSS

* Use CCSS context also for cloud and move retried state to it

* Rename variable for clarity

* Update retried state after regeneration has started
---
 .../cloud-css-meta/cloud-css-meta.tsx         |  4 --
 .../critical-css-context-provider.tsx}        | 40 ++++++++++++++-----
 .../critical-css-meta/critical-css-meta.tsx   |  8 +---
 .../critical-css/lib/critical-css-errors.ts   |  5 +++
 .../lib/stores/critical-css-state.ts          | 38 +++++++++++++++---
 .../critical-css/lib/use-retry-regenerate.ts  |  9 ++---
 .../show-stopper-error/show-stopper-error.tsx | 20 ++++------
 .../features/critical-css/status/status.tsx   |  6 ---
 .../js/features/speed-score/speed-score.tsx   |  2 +-
 .../js/layout/settings-page/settings-page.tsx |  6 +--
 .../boost/changelog/fix-critical-css-status   |  4 ++
 11 files changed, 87 insertions(+), 55 deletions(-)
 rename projects/plugins/boost/app/assets/src/js/features/critical-css/{local-generator/local-generator-provider.tsx => critical-css-context/critical-css-context-provider.tsx} (73%)
 create mode 100644 projects/plugins/boost/changelog/fix-critical-css-status

diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/cloud-css-meta/cloud-css-meta.tsx b/projects/plugins/boost/app/assets/src/js/features/critical-css/cloud-css-meta/cloud-css-meta.tsx
index 8f1f3ae65660c..7dc510ff0e5ce 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/cloud-css-meta/cloud-css-meta.tsx
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/cloud-css-meta/cloud-css-meta.tsx
@@ -1,12 +1,10 @@
 import { __ } from '@wordpress/i18n';
 import Status from '../status/status';
 import { useCriticalCssState } from '../lib/stores/critical-css-state';
-import { useRetryRegenerate } from '../lib/use-retry-regenerate';
 import { isFatalError } from '../lib/critical-css-errors';
 
 export default function CloudCssMetaProps() {
 	const [ cssState ] = useCriticalCssState();
-	const [ hasRetried, retry ] = useRetryRegenerate();
 
 	const isPending = cssState.status === 'pending';
 	const hasCompletedSome = cssState.providers.some( provider => provider.status !== 'pending' );
@@ -29,8 +27,6 @@ export default function CloudCssMetaProps() {
 			cssState={ cssState }
 			isCloud={ true }
 			showFatalError={ isFatalError( cssState ) }
-			hasRetried={ hasRetried }
-			retry={ retry }
 			extraText={ extraText || undefined }
 			overrideText={ overrideText || undefined }
 		/>
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/local-generator/local-generator-provider.tsx b/projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-context/critical-css-context-provider.tsx
similarity index 73%
rename from projects/plugins/boost/app/assets/src/js/features/critical-css/local-generator/local-generator-provider.tsx
rename to projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-context/critical-css-context-provider.tsx
index deb0e59881b6e..0b2d52bffddfd 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/local-generator/local-generator-provider.tsx
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-context/critical-css-context-provider.tsx
@@ -11,19 +11,23 @@ import {
 import { runLocalGenerator } from '../lib/generate-critical-css';
 import { CriticalCssErrorDetails } from '../lib/stores/critical-css-state-types';
 
-type LocalGeneratorContext = {
+type CriticalCssContextValues = {
 	isGenerating: boolean;
 	setGenerating: ( generating: boolean ) => void;
 
 	providerProgress: number;
 	setProviderProgress: ( progress: number ) => void;
+
+	// Whether we've retried generating critical CSS after an error.
+	hasRetriedAfterError: boolean;
+	setHasRetriedAfterError: ( hasRetried: boolean ) => void;
 };
 
 type ProviderProps = {
 	children: ReactNode;
 };
 
-const CssGeneratorContext = createContext< LocalGeneratorContext | null >( null );
+const CriticalCssContext = createContext< CriticalCssContextValues | null >( null );
 
 /**
  * Local Critical CSS Context Provider component - provides context for any descendants that want to
@@ -31,28 +35,34 @@ const CssGeneratorContext = createContext< LocalGeneratorContext | null >( null
  *
  * @param {ProviderProps} props - Component props.
  */
-export default function LocalCriticalCssGeneratorProvider( { children }: ProviderProps ) {
+export default function CriticalCssProvider( { children }: ProviderProps ) {
 	const [ isGenerating, setGenerating ] = useState< boolean >( false );
 	const [ providerProgress, setProviderProgress ] = useState< number >( 0 );
+	const [ hasRetriedAfterError, setHasRetriedAfterError ] = useState< boolean >( false );
 
 	const value = {
+		// Local Generator status.
 		isGenerating,
 		setGenerating,
 		providerProgress,
 		setProviderProgress,
+
+		// Whether we've retried generating critical CSS after an error.
+		hasRetriedAfterError,
+		setHasRetriedAfterError,
 	};
 
-	return <CssGeneratorContext.Provider value={ value }>{ children }</CssGeneratorContext.Provider>;
+	return <CriticalCssContext.Provider value={ value }>{ children }</CriticalCssContext.Provider>;
 }
 
 /**
  * Internal helper function: Use the raw Critical CSS Generator context, and verify it's inside a provider.
  */
-function useLocalCriticalCssGeneratorContext() {
-	const status = useContext( CssGeneratorContext );
+function useCriticalCssContext() {
+	const status = useContext( CriticalCssContext );
 
 	if ( ! status ) {
-		throw new Error( 'Local critical CSS generator status not available' );
+		throw new Error( 'Critical CSS status not available' );
 	}
 
 	return status;
@@ -62,18 +72,26 @@ function useLocalCriticalCssGeneratorContext() {
  * For status consumers: Get an overview of the local critical CSS generator status. Is it running or not?
  */
 export function useLocalCriticalCssGeneratorStatus() {
-	const { isGenerating, providerProgress } = useLocalCriticalCssGeneratorContext();
+	const { isGenerating, providerProgress } = useCriticalCssContext();
 
 	return { isGenerating, providerProgress };
 }
 
+/** The retried state of critical CSS. */
+export function useCriticalCssRetriedAfterErrorState() {
+	const { hasRetriedAfterError: hasRetried, setHasRetriedAfterError: setHasRetried } =
+		useCriticalCssContext();
+
+	return [ hasRetried, setHasRetried ] as const;
+}
+
 /**
  * For Critical CSS UI: Actually run the local generator and return its status.
  */
 export function useLocalCriticalCssGenerator() {
 	// Local Generator status context.
 	const { isGenerating, setGenerating, providerProgress, setProviderProgress } =
-		useLocalCriticalCssGeneratorContext();
+		useCriticalCssContext();
 
 	// Critical CSS state and actions.
 	const [ cssState, setCssState ] = useCriticalCssState();
@@ -86,7 +104,7 @@ export function useLocalCriticalCssGenerator() {
 
 	useEffect(
 		() => {
-			if ( cssState.status === 'pending' ) {
+			if ( cssState.status === 'pending' && cssState.providers.length > 0 ) {
 				let abortController: AbortController | undefined;
 
 				setGenerating( true );
@@ -119,7 +137,7 @@ export function useLocalCriticalCssGenerator() {
 		// This effect triggers an actual process that is costly to start and stop, so we don't want to start/stop it
 		// every time an object ref like `cssState` is changed for a trivial reason.
 		// eslint-disable-next-line react-hooks/exhaustive-deps
-		[ cssState.status ]
+		[ cssState.status, cssState.providers.length ]
 	);
 
 	const progress =
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-meta/critical-css-meta.tsx b/projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-meta/critical-css-meta.tsx
index c55cdc1532b17..bde75368908b5 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-meta/critical-css-meta.tsx
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/critical-css-meta/critical-css-meta.tsx
@@ -4,8 +4,7 @@ import ProgressBar from '$features/ui/progress-bar/progress-bar';
 import styles from './critical-css-meta.module.scss';
 import { useCriticalCssState } from '../lib/stores/critical-css-state';
 import { RegenerateCriticalCssSuggestion, useRegenerationReason } from '..';
-import { useLocalCriticalCssGenerator } from '../local-generator/local-generator-provider';
-import { useRetryRegenerate } from '../lib/use-retry-regenerate';
+import { useLocalCriticalCssGenerator } from '../critical-css-context/critical-css-context-provider';
 import { isFatalError } from '../lib/critical-css-errors';
 
 /**
@@ -14,12 +13,11 @@ import { isFatalError } from '../lib/critical-css-errors';
  */
 export default function CriticalCssMeta() {
 	const [ cssState ] = useCriticalCssState();
-	const [ hasRetried, retry ] = useRetryRegenerate();
 	const [ { data: regenerateReason } ] = useRegenerationReason();
 	const { progress } = useLocalCriticalCssGenerator();
 	const showFatalError = isFatalError( cssState );
 
-	if ( cssState.status === 'pending' ) {
+	if ( cssState.status === 'pending' || cssState.status === 'not_generated' ) {
 		return (
 			<div className="jb-critical-css-progress">
 				<div className={ styles[ 'progress-label' ] }>
@@ -39,8 +37,6 @@ export default function CriticalCssMeta() {
 				cssState={ cssState }
 				isCloud={ false }
 				showFatalError={ showFatalError }
-				hasRetried={ hasRetried }
-				retry={ retry }
 				highlightRegenerateButton={ !! regenerateReason }
 				extraText={ __(
 					'Remember to regenerate each time you make changes that affect your HTML or CSS structure.',
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/critical-css-errors.ts b/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/critical-css-errors.ts
index 52260adb4b0b3..7967660eead18 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/critical-css-errors.ts
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/critical-css-errors.ts
@@ -34,6 +34,11 @@ export function isFatalError( cssState: CriticalCssState ): boolean {
 		return false;
 	}
 
+	// If there are no providers, the state is being re-initialized. So dismiss any show-stopper errors.
+	if ( cssState.providers.length === 0 ) {
+		return false;
+	}
+
 	const hasActiveProvider = cssState.providers.some(
 		provider => provider.status === 'success' || provider.status === 'pending'
 	);
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/stores/critical-css-state.ts b/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/stores/critical-css-state.ts
index ec0b282d2250b..b5e7da70af7b7 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/stores/critical-css-state.ts
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/stores/critical-css-state.ts
@@ -57,14 +57,20 @@ export function criticalCssErrorState( message: string ): CriticalCssState {
  * All Critical CSS State actions return a success flag and the new state. This hook wraps the
  * common logic for handling the result of these actions.
  *
- * @param {string}      action    - The name of the action.
- * @param {z.ZodSchema} schema    - The schema for the action request.
- * @param {Function}    onSuccess - Optional callback for handling the new state.
+ * @param {string}           action          - The name of the action.
+ * @param {z.ZodSchema}      schema          - The schema for the action request.
+ * @param {CriticalCssState} optimisticState - The state to use for optimistic updates.
+ * @param {Function}         onSuccess       - Optional callback for handling the new state.
  */
 function useCriticalCssAction<
 	ActionSchema extends z.ZodSchema,
 	ActionRequestData extends z.infer< ActionSchema >,
->( action: string, schema: ActionRequestData, onSuccess?: ( state: CriticalCssState ) => void ) {
+>(
+	action: string,
+	schema: ActionRequestData,
+	optimisticState?: CriticalCssState,
+	onSuccess?: ( state: CriticalCssState ) => void
+) {
 	const responseSchema = z.object( {
 		success: z.boolean(),
 		state: CriticalCssStateSchema,
@@ -90,6 +96,13 @@ function useCriticalCssAction<
 			action_response: responseSchema,
 		},
 		callbacks: {
+			optimisticUpdate: ( _requestData, state: CriticalCssState ) => {
+				if ( optimisticState ) {
+					return optimisticState;
+				}
+
+				return state;
+			},
 			onResult: ( result, _state ): CriticalCssState => {
 				if ( result.success ) {
 					if ( onSuccess ) {
@@ -150,10 +163,23 @@ export function useSetProviderErrorsAction() {
 
 /**
  * Hook which creates a callable action for regenerating Critical CSS.
+ *
+ * @param {Function} callback - Optional callback to call when a regeneration starts successfully.
  */
-export function useRegenerateCriticalCssAction() {
+export function useRegenerateCriticalCssAction( callback?: () => void ) {
 	const [ , resetReason ] = useRegenerationReason();
-	return useCriticalCssAction( 'request-regenerate', z.void(), resetReason );
+
+	const onSuccess = () => {
+		if ( callback ) {
+			callback();
+		}
+
+		resetReason();
+	};
+
+	// Optimistically update the state to hide any errors and immediately show the pending state.
+	const optimisticState: CriticalCssState = { status: 'pending', providers: [] };
+	return useCriticalCssAction( 'request-regenerate', z.void(), optimisticState, onSuccess );
 }
 
 /**
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/use-retry-regenerate.ts b/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/use-retry-regenerate.ts
index c8ef9b5bf6a97..01e4b193271c1 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/use-retry-regenerate.ts
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/lib/use-retry-regenerate.ts
@@ -1,5 +1,5 @@
-import { useState } from 'react';
 import { useRegenerateCriticalCssAction } from './stores/critical-css-state';
+import { useCriticalCssRetriedAfterErrorState } from '../critical-css-context/critical-css-context-provider';
 
 /**
  * Helper for "Retry" buttons for Critical CSS which need to track whether they have been clicked
@@ -8,13 +8,12 @@ import { useRegenerateCriticalCssAction } from './stores/critical-css-state';
  * Returns a boolean indicating whether retrying has been attempted, and a function to call to retry.
  */
 export function useRetryRegenerate(): [ boolean, () => void ] {
-	const [ retried, setRetried ] = useState( false );
-	const regenerateAction = useRegenerateCriticalCssAction();
+	const [ retriedAfterError, setRetriedAfterError ] = useCriticalCssRetriedAfterErrorState();
+	const regenerateAction = useRegenerateCriticalCssAction( () => setRetriedAfterError( true ) );
 
 	function retry() {
-		setRetried( true );
 		regenerateAction.mutate();
 	}
 
-	return [ retried, retry ];
+	return [ retriedAfterError, retry ];
 }
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/show-stopper-error/show-stopper-error.tsx b/projects/plugins/boost/app/assets/src/js/features/critical-css/show-stopper-error/show-stopper-error.tsx
index 5e640257c58f6..e46c527cd8a3d 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/show-stopper-error/show-stopper-error.tsx
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/show-stopper-error/show-stopper-error.tsx
@@ -11,19 +11,16 @@ import getCriticalCssErrorSetInterpolateVars from '$lib/utils/get-critical-css-e
 import formatErrorSetUrls from '$lib/utils/format-error-set-urls';
 import actionLinkInterpolateVar from '$lib/utils/action-link-interpolate-var';
 import { recordBoostEvent } from '$lib/utils/analytics';
+import { useRetryRegenerate } from '../lib/use-retry-regenerate';
 
 type ShowStopperErrorTypes = {
 	supportLink?: string;
 	cssState: CriticalCssState;
-	retry: () => void;
-	showRetry?: boolean;
 };
 
 const ShowStopperError: React.FC< ShowStopperErrorTypes > = ( {
 	supportLink = 'https://wordpress.org/support/plugin/jetpack-boost/',
 	cssState,
-	retry,
-	showRetry,
 } ) => {
 	const primaryErrorSet = getPrimaryErrorSet( cssState );
 	const showLearnSection = primaryErrorSet && cssState.status === 'generated';
@@ -55,12 +52,7 @@ const ShowStopperError: React.FC< ShowStopperErrorTypes > = ( {
 						</FoldingElement>
 					</>
 				) : (
-					<OtherErrors
-						cssState={ cssState }
-						retry={ retry }
-						showRetry={ showRetry }
-						supportLink={ supportLink }
-					/>
+					<OtherErrors cssState={ cssState } supportLink={ supportLink } />
 				) }
 			</Notice>
 		</>
@@ -136,7 +128,9 @@ const DocumentationSection = ( {
 	);
 };
 
-const OtherErrors = ( { cssState, retry, showRetry, supportLink }: ShowStopperErrorTypes ) => {
+const OtherErrors = ( { cssState, supportLink }: ShowStopperErrorTypes ) => {
+	const [ hasRetried, retry ] = useRetryRegenerate();
+
 	const firstTimeError = __(
 		'An unexpected error has occurred. As this error may be temporary, please try and refresh the Critical CSS.',
 		'jetpack-boost'
@@ -184,7 +178,7 @@ const OtherErrors = ( { cssState, retry, showRetry, supportLink }: ShowStopperEr
 				</>
 			) : (
 				<>
-					<p>{ showRetry ? firstTimeError : secondTimeError }</p>
+					<p>{ ! hasRetried ? firstTimeError : secondTimeError }</p>
 					<p>
 						{ sprintf(
 							/* translators: %s: error message */
@@ -192,7 +186,7 @@ const OtherErrors = ( { cssState, retry, showRetry, supportLink }: ShowStopperEr
 							cssState.status_error
 						) }
 					</p>
-					{ showRetry ? (
+					{ ! hasRetried ? (
 						<button
 							className="secondary"
 							onClick={ () => {
diff --git a/projects/plugins/boost/app/assets/src/js/features/critical-css/status/status.tsx b/projects/plugins/boost/app/assets/src/js/features/critical-css/status/status.tsx
index d3b7d596ced7f..292e1d2643c54 100644
--- a/projects/plugins/boost/app/assets/src/js/features/critical-css/status/status.tsx
+++ b/projects/plugins/boost/app/assets/src/js/features/critical-css/status/status.tsx
@@ -17,8 +17,6 @@ type StatusTypes = {
 	cssState: CriticalCssState;
 	isCloud?: boolean;
 	showFatalError: boolean;
-	hasRetried: boolean;
-	retry: () => void;
 	highlightRegenerateButton?: boolean;
 	extraText?: string; // Optionally, provide a sentence to use after the main message to provide more context.
 	overrideText?: string; // Optionally, provide a custom message to display instead of the default.
@@ -28,8 +26,6 @@ const Status: React.FC< StatusTypes > = ( {
 	cssState,
 	isCloud = false,
 	showFatalError,
-	hasRetried,
-	retry,
 	highlightRegenerateButton = false,
 	extraText,
 	overrideText,
@@ -54,8 +50,6 @@ const Status: React.FC< StatusTypes > = ( {
 			<ShowStopperError
 				supportLink={ ( isCloud && 'https://jetpack.com/contact-support/' ) || undefined }
 				cssState={ cssState }
-				retry={ retry }
-				showRetry={ ! hasRetried }
 			/>
 		);
 	}
diff --git a/projects/plugins/boost/app/assets/src/js/features/speed-score/speed-score.tsx b/projects/plugins/boost/app/assets/src/js/features/speed-score/speed-score.tsx
index 7010f0af185fb..127f9f3cfe176 100644
--- a/projects/plugins/boost/app/assets/src/js/features/speed-score/speed-score.tsx
+++ b/projects/plugins/boost/app/assets/src/js/features/speed-score/speed-score.tsx
@@ -16,7 +16,7 @@ import { useDebouncedRefreshScore, useSpeedScores } from './lib/hooks';
 import styles from './speed-score.module.scss';
 import { useModulesState } from '$features/module/lib/stores';
 import { useCriticalCssState } from '$features/critical-css/lib/stores/critical-css-state';
-import { useLocalCriticalCssGeneratorStatus } from '$features/critical-css/local-generator/local-generator-provider';
+import { useLocalCriticalCssGeneratorStatus } from '$features/critical-css/critical-css-context/critical-css-context-provider';
 import { queryClient } from '@automattic/jetpack-react-data-sync-client';
 import ErrorBoundary from '$features/error-boundary/error-boundary';
 import PopOut from './pop-out/pop-out';
diff --git a/projects/plugins/boost/app/assets/src/js/layout/settings-page/settings-page.tsx b/projects/plugins/boost/app/assets/src/js/layout/settings-page/settings-page.tsx
index 7162492dd54eb..7b86734da8f41 100644
--- a/projects/plugins/boost/app/assets/src/js/layout/settings-page/settings-page.tsx
+++ b/projects/plugins/boost/app/assets/src/js/layout/settings-page/settings-page.tsx
@@ -6,7 +6,7 @@ import Tips from './tips/tips';
 import clsx from 'clsx';
 import styles from './settings-page.module.scss';
 import { usePremiumFeatures } from '$lib/stores/premium-features';
-import LocalCriticalCssGeneratorProvider from '$features/critical-css/local-generator/local-generator-provider';
+import CriticalCssProvider from '$features/critical-css/critical-css-context/critical-css-context-provider';
 import NoticeManager from '$features/notice/manager';
 import { NoticeProvider } from '$features/notice/context';
 
@@ -20,7 +20,7 @@ const SettingsPage = ( { children }: SettingsPageProps ) => {
 
 	return (
 		<NoticeProvider>
-			<LocalCriticalCssGeneratorProvider>
+			<CriticalCssProvider>
 				<div id="jb-dashboard" className="jb-dashboard jb-dashboard--main">
 					<Header />
 
@@ -41,7 +41,7 @@ const SettingsPage = ( { children }: SettingsPageProps ) => {
 					<Footer />
 					<NoticeManager />
 				</div>
-			</LocalCriticalCssGeneratorProvider>
+			</CriticalCssProvider>
 		</NoticeProvider>
 	);
 };
diff --git a/projects/plugins/boost/changelog/fix-critical-css-status b/projects/plugins/boost/changelog/fix-critical-css-status
new file mode 100644
index 0000000000000..fd96c0c4e40a4
--- /dev/null
+++ b/projects/plugins/boost/changelog/fix-critical-css-status
@@ -0,0 +1,4 @@
+Significance: patch
+Type: fixed
+
+Critical CSS: Improved UI responsiveness during a retry after failed generation.