From 3bc2f1a740686bcd6db2f3ea9c3fa94fe2b79ecd Mon Sep 17 00:00:00 2001
From: Mati Horowitz <mati@elementor.com>
Date: Tue, 14 Nov 2023 20:17:20 +0200
Subject: [PATCH] Fix: AI - Screenshots are not being generated properly

---
 modules/ai/assets/js/editor/utils/screenshot.js | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/modules/ai/assets/js/editor/utils/screenshot.js b/modules/ai/assets/js/editor/utils/screenshot.js
index b486593c4b4d..04d0053bd840 100644
--- a/modules/ai/assets/js/editor/utils/screenshot.js
+++ b/modules/ai/assets/js/editor/utils/screenshot.js
@@ -20,6 +20,10 @@ export const takeScreenshot = async ( template ) => {
 	// Wait for the container to render.
 	await waitForContainer( container.id );
 
+	if ( template.elements.length ) {
+		await Promise.all( template.elements.map( ( child ) => waitForContainer( child.id ) ) );
+	}
+
 	let screenshot;
 
 	try {