From fb9d637681b740d1f0b580940a54857ad61634e1 Mon Sep 17 00:00:00 2001
From: Douglas Henri <dhasilva@users.noreply.github.com>
Date: Wed, 30 Aug 2023 14:54:41 -0300
Subject: [PATCH] AI Extension: Enhance blocks parsing flickering with child
 blocks (#32635)

* add block comparison function

* add prompt tweaks and history with examples

* fix streaming flickering

* changelog

* update changelog
---
 .../update-ai-assistant-required-fields       |  4 +++
 .../with-ui-handler-data-provider.tsx         | 25 +++++++++++--
 .../ai-assistant/lib/utils/compare-blocks.ts  | 35 +++++++++++++++++++
 3 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 projects/plugins/jetpack/changelog/update-ai-assistant-required-fields
 create mode 100644 projects/plugins/jetpack/extensions/blocks/ai-assistant/lib/utils/compare-blocks.ts

diff --git a/projects/plugins/jetpack/changelog/update-ai-assistant-required-fields b/projects/plugins/jetpack/changelog/update-ai-assistant-required-fields
new file mode 100644
index 0000000000000..568472ceb4c5e
--- /dev/null
+++ b/projects/plugins/jetpack/changelog/update-ai-assistant-required-fields
@@ -0,0 +1,4 @@
+Significance: patch
+Type: enhancement
+
+AI Extension: Enhance blocks parsing flickering with child blocks
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx
index e1d27f5b1d01c..2aa1c425f5026 100644
--- a/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/extensions/jetpack-contact-form/ui-handler/with-ui-handler-data-provider.tsx
@@ -14,6 +14,7 @@ import { store as noticesStore } from '@wordpress/notices';
  * Internal dependencies
  */
 import { isPossibleToExtendJetpackFormBlock } from '..';
+import { compareBlocks } from '../../../lib/utils/compare-blocks';
 import { fixIncompleteHTML } from '../../../lib/utils/fix-incomplete-html';
 import { AiAssistantUiContextProvider } from './context';
 /**
@@ -167,8 +168,28 @@ const withUiHandlerDataProvider = createHigherOrderComponent( BlockListBlock =>
 					return block.isValid && block.name !== 'core/freeform' && block.name !== 'core/missing';
 				} );
 
-				// Only update the blocks when the valid list changed, meaning a new block arrived.
-				if ( validBlocks.length !== currentListOfValidBlocks.current.length ) {
+				let lastBlockUpdated = false;
+
+				// While streaming, the last block can go from valid to invalid and back as new children are added token by token.
+				if ( validBlocks.length < currentListOfValidBlocks.current.length ) {
+					// The last block is temporarily invalid, so we use the last valid state.
+					validBlocks.push(
+						currentListOfValidBlocks.current[ currentListOfValidBlocks.current.length - 1 ]
+					);
+				} else if (
+					validBlocks.length === currentListOfValidBlocks.current.length &&
+					validBlocks.length > 0
+				) {
+					// Update the last valid block with the new content if it is different.
+					const lastBlock = validBlocks[ validBlocks.length - 1 ];
+					const lastBlockFromCurrentList =
+						currentListOfValidBlocks.current[ validBlocks.length - 1 ];
+
+					lastBlockUpdated = ! compareBlocks( lastBlock, lastBlockFromCurrentList );
+				}
+
+				// Only update the blocks when the valid list changed, meaning a new block arrived or the last block was updated.
+				if ( validBlocks.length !== currentListOfValidBlocks.current.length || lastBlockUpdated ) {
 					// Only update the valid blocks
 					replaceInnerBlocks( clientId, validBlocks );
 
diff --git a/projects/plugins/jetpack/extensions/blocks/ai-assistant/lib/utils/compare-blocks.ts b/projects/plugins/jetpack/extensions/blocks/ai-assistant/lib/utils/compare-blocks.ts
new file mode 100644
index 0000000000000..9e50257fecee1
--- /dev/null
+++ b/projects/plugins/jetpack/extensions/blocks/ai-assistant/lib/utils/compare-blocks.ts
@@ -0,0 +1,35 @@
+type Block = {
+	attributes?: object;
+	clientId?: string;
+	innerBlocks?: Block[];
+	isValid?: boolean;
+	name?: string;
+	originalContent?: string;
+};
+
+const omitClientId = ( block: Block ): Block => {
+	delete block.clientId;
+
+	for ( const child of block.innerBlocks ?? [] ) {
+		omitClientId( child );
+	}
+
+	return block;
+};
+
+const copyBlock = ( block: Block ): Block => JSON.parse( JSON.stringify( block ) );
+const copyBlockWithoutClientId = ( block: Block ) => omitClientId( copyBlock( block ) );
+
+/**
+ * Deeply compares two blocks, ignoring the clientId property.
+ *
+ * @param {Block} blockA - The first block to compare.
+ * @param {Block} blockB - The second block to compare.
+ * @returns {boolean} Whether the two blocks are equal.
+ */
+export function compareBlocks( blockA: Block, blockB: Block ): boolean {
+	const aCopy = copyBlockWithoutClientId( blockA );
+	const bCopy = copyBlockWithoutClientId( blockB );
+
+	return JSON.stringify( aCopy ) === JSON.stringify( bCopy );
+}