Restore classification suggestion pill to original position on API failure

src/experiments/content-classification/components/useContentClassification.ts

@@ -34,6 +34,14 @@ const MAX_SUGGESTIONS = 10;
 
 const suggestionsCache: Record< string, TagSuggestion[] > = {};
 
+/**
+ * Tracks the current suggestions generation cycle count for each taxonomy.
+ *
+ * Incremented on suggestion regeneration or dismissal.
+ * Prevents race conditions from stale request failures trying to restore pills.
+ */
+const generationCounts: Record< string, number > = {};
+
 const normalizeMaxSuggestions = ( value: unknown ): number => {
 	const parsedValue = Number.parseInt(
 		String( value ?? DEFAULT_MAX_SUGGESTIONS ),
@@ -178,6 +186,9 @@ export function useContentClassification( taxonomy: string ): {
 			return;
 		}
 
+		generationCounts[ taxonomy ] =
+			( generationCounts[ taxonomy ] || 0 ) + 1;
+
 		const settings = getSettings();
 		setIsGenerating( true );
 		setSuggestions( [] );
@@ -226,8 +237,32 @@ export function useContentClassification( taxonomy: string ): {
 	// Handle accepting a suggestion.
 	const handleAccept = useCallback(
 		( suggestion: TagSuggestion ) => {
-			removeSuggestionFromList( suggestion );
-			addTermToPost( taxonomy, suggestion );
+			// Obtain original position of the suggestion before removal.
+			let originalIndex = -1;
+			setSuggestions( ( prev ) => {
+				originalIndex = prev.findIndex(
+					( s ) => s.term === suggestion.term
+				);
+				return prev.filter( ( s ) => s.term !== suggestion.term );
+			} );
+
+			const clickGenerationCount = generationCounts[ taxonomy ] || 0;
+
+			addTermToPost( taxonomy, suggestion ).then( ( success ) => {
+				if ( success ) {
+					return;
+				}
+
+				// Ensure to not add suggestion back if a new generation/clear has been triggered.
+				if (
+					( generationCounts[ taxonomy ] || 0 ) ===
+					clickGenerationCount
+				) {
+					setSuggestions( ( prev ) =>
+						insertSuggestionAt( prev, suggestion, originalIndex )
+					);
+				}
+			} );
 		},
 		[ taxonomy ]
 	);
@@ -239,8 +274,10 @@ export function useContentClassification( taxonomy: string ): {
 
 	// Handle dismissing all suggestions.
 	const handleDismissAll = useCallback( () => {
+		generationCounts[ taxonomy ] =
+			( generationCounts[ taxonomy ] || 0 ) + 1;
 		setSuggestions( [] );
-	}, [] );
+	}, [ taxonomy ] );
 
 	return {
 		isGenerating,
@@ -254,19 +291,46 @@ export function useContentClassification( taxonomy: string ): {
 	};
 }
 
+/**
+ * Inserts a suggestion back into a list at the specified index, preventing duplicates.
+ *
+ * @param list       The current list of suggestions.
+ * @param suggestion The suggestion to insert.
+ * @param index      The index to insert the suggestion at.
+ * @return The updated list of suggestions.
+ */
+function insertSuggestionAt(
+	list: TagSuggestion[],
+	suggestion: TagSuggestion,
+	index: number
+): TagSuggestion[] {
+	// Skip if already present.
+	if ( list.some( ( s ) => s.term === suggestion.term ) ) {
+		return list;
+	}
+
+	// If original index is invalid, push to the end.
+	if ( index < 0 || index > list.length ) {
+		return [ ...list, suggestion ];
+	}
+
+	// Otherwise, insert at the correct position.
+	const newList = [ ...list ];
+	newList.splice( index, 0, suggestion );
+	return newList;
+}
+
 /**
  * Adds a term to the current post.
  *
  * @param taxonomy   The taxonomy slug.
  * @param suggestion The suggestion to add.
+ * @return Promise resolving to true if successful, false otherwise.
  */
 async function addTermToPost(
 	taxonomy: string,
 	suggestion: TagSuggestion
-): Promise< void > {
-	const { editPost }: any = dispatch( editorStore );
-	const { getEditedPostAttribute } = select( editorStore );
-
+): Promise< boolean > {
 	const taxonomyObject: any = select( coreStore ).getTaxonomy( taxonomy );
 	const restBase = taxonomyObject?.rest_base ?? taxonomy;
 
@@ -280,9 +344,14 @@ async function addTermToPost(
 		);
 		if ( resolvedParent ) {
 			parentId = resolvedParent;
+		} else {
+			return false;
 		}
 	}
 
+	const { editPost }: any = dispatch( editorStore );
+	const { getEditedPostAttribute } = select( editorStore );
+
 	const currentTerms: number[] = getEditedPostAttribute( restBase ) ?? [];
 	const termId = await findOrCreateTerm(
 		taxonomy,
@@ -291,11 +360,16 @@ async function addTermToPost(
 		parentId
 	);
 
-	if ( termId && ! currentTerms.includes( termId ) ) {
-		editPost( {
-			[ restBase ]: [ ...currentTerms, termId ],
-		} );
+	if ( termId ) {
+		if ( ! currentTerms.includes( termId ) ) {
+			editPost( {
+				[ restBase ]: [ ...currentTerms, termId ],
+			} );
+		}
+		return true;
 	}
+
+	return false;
 }
 
 /**

tests/e2e/specs/experiments/content-classification.spec.js

@@ -460,4 +460,158 @@ test.describe( 'Content Classification Experiment', () => {
 			page.locator( '.ai-content-classification__pill' ).first()
 		).toBeVisible();
 	} );
+
+	test( 'Restores suggestion pill to original position if tag addition fails', async ( {
+		admin,
+		editor,
+		page,
+	} ) => {
+		await setStrategy( admin, page, 'allow_new' );
+
+		await admin.createNewPost( {
+			title: 'Content Classification Pill Restoration Failure Test',
+		} );
+
+		await editor.insertBlock( {
+			name: 'core/paragraph',
+			attributes: { content: LONG_CONTENT },
+		} );
+
+		// Open the Tags panel and suggest.
+		await openTaxonomyPanel( editor, page, 'Tags' );
+		await page
+			.locator( '.ai-content-classification button', {
+				hasText: 'Suggest Tags',
+			} )
+			.first()
+			.click();
+
+		await expect(
+			page.locator( '.ai-content-classification__pill' ).first()
+		).toBeVisible();
+
+		let resolveRoute;
+		const routePromise = new Promise( ( resolve ) => {
+			resolveRoute = resolve;
+		} );
+
+		// Intercept tag creation/search REST calls to fail.
+		await page.route( /\/wp\/v2\/tags/, async ( route ) => {
+			await routePromise;
+			await route.fulfill( {
+				status: 500,
+				contentType: 'application/json',
+				body: JSON.stringify( {
+					code: 'internal_server_error',
+					message: 'Database connection failed.',
+				} ),
+			} );
+		} );
+
+		const firstPill = page
+			.locator( '.ai-content-classification__pill' )
+			.first();
+		const pillText = await firstPill
+			.locator( '.ai-content-classification__pill-accept' )
+			.textContent();
+
+		// Click to accept a term.
+		await firstPill
+			.locator( '.ai-content-classification__pill-accept' )
+			.click();
+
+		// Pill should disappear immediately.
+		await expect(
+			page.locator( '.ai-content-classification__pill-accept' ).first()
+		).not.toHaveText( pillText );
+
+		// Resolve the promise to return the failure response.
+		resolveRoute();
+
+		// Since request fails, the pill should be restored to its original (first) position.
+		await expect(
+			page.locator( '.ai-content-classification__pill-accept' ).first()
+		).toHaveText( pillText );
+	} );
+
+	test( 'Does not restore suggestion pill if it belongs to a stale generation session', async ( {
+		admin,
+		editor,
+		page,
+	} ) => {
+		await setStrategy( admin, page, 'allow_new' );
+
+		await admin.createNewPost( {
+			title: 'Content Classification Stale Generation Test',
+		} );
+
+		await editor.insertBlock( {
+			name: 'core/paragraph',
+			attributes: { content: LONG_CONTENT },
+		} );
+
+		// Open the Tags panel and suggest.
+		await openTaxonomyPanel( editor, page, 'Tags' );
+		await page
+			.locator( '.ai-content-classification button', {
+				hasText: 'Suggest Tags',
+			} )
+			.first()
+			.click();
+
+		await expect(
+			page.locator( '.ai-content-classification__pill' ).first()
+		).toBeVisible();
+
+		let resolveRoute;
+		const routePromise = new Promise( ( resolve ) => {
+			resolveRoute = resolve;
+		} );
+
+		// Intercept tag creation/search REST calls to fail.
+		await page.route( /\/wp\/v2\/tags/, async ( route ) => {
+			await routePromise;
+			await route.fulfill( {
+				status: 500,
+				contentType: 'application/json',
+				body: JSON.stringify( {
+					code: 'internal_server_error',
+					message: 'Database connection failed.',
+				} ),
+			} );
+		} );
+
+		const firstPill = page
+			.locator( '.ai-content-classification__pill' )
+			.first();
+		const pillText = await firstPill
+			.locator( '.ai-content-classification__pill-accept' )
+			.textContent();
+
+		// Click to accept a term.
+		await firstPill
+			.locator( '.ai-content-classification__pill-accept' )
+			.click();
+
+		// Pill should disappear immediately.
+		await expect(
+			page.locator( '.ai-content-classification__pill-accept' ).first()
+		).not.toHaveText( pillText );
+
+		// While API call is pending, click "Dismiss all" to clear suggestions and increment session count.
+		await page
+			.locator( '.ai-content-classification__actions button', {
+				hasText: 'Dismiss all',
+			} )
+			.first()
+			.click();
+
+		// Resolve the promise to return the failure response.
+		resolveRoute();
+
+		// Since the session is now stale, the suggestions UI should remain cleared/empty.
+		await expect(
+			page.locator( '.ai-content-classification__suggestions' ).first()
+		).not.toBeVisible();
+	} );
 } );