From 60fbd16730c72cff816e687487e2ef331bb97bfc Mon Sep 17 00:00:00 2001
From: magmueller <mamagnus00@gmail.com>
Date: Sat, 23 Nov 2024 15:39:47 +0100
Subject: [PATCH] When page not loaded return empty state - fixes break for not
 loading page

---
 .gitignore                                   |  4 +++-
 browser_use/agent/service.py                 | 14 +++++++-------
 browser_use/browser/tests/playwright_test.py |  2 ++
 browser_use/dom/service.py                   |  1 +
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9bc4d79bb..4bdae02ff 100644
--- a/.gitignore
+++ b/.gitignore
@@ -164,4 +164,6 @@ temp
 tmp
 
 
-.DS_Store
\ No newline at end of file
+.DS_Store
+
+private_example.py
\ No newline at end of file
diff --git a/browser_use/agent/service.py b/browser_use/agent/service.py
index 9dec94c29..98ed90a56 100644
--- a/browser_use/agent/service.py
+++ b/browser_use/agent/service.py
@@ -117,9 +117,9 @@ class Agent:
 	async def step(self) -> None:
 		"""Execute one step of the task"""
 		logger.info(f'\n📍 Step {self.n_steps}')
-		state = await self.controller.browser.get_state(use_vision=self.use_vision)
-
+		state = None
 		try:
+			state = await self.controller.browser.get_state(use_vision=self.use_vision)
 			model_output = await self.get_next_action(state)
 			result = await self.controller.act(model_output.action)
 			if result.extracted_content:
@@ -129,7 +129,7 @@ class Agent:
 			self.consecutive_failures = 0
 
 		except Exception as e:
-			result = self._handle_step_error(e, state)
+			result = self._handle_step_error(e)
 			model_output = None
 
 			if result.error:
@@ -139,11 +139,11 @@ class Agent:
 						error=result.error,
 					)
 				)
+		if state:
+			self._update_messages_with_result(result)
+			self._make_history_item(model_output, state, result)
 
-		self._update_messages_with_result(result)
-		self._make_history_item(model_output, state, result)
-
-	def _handle_step_error(self, error: Exception, state: BrowserState) -> ActionResult:
+	def _handle_step_error(self, error: Exception) -> ActionResult:
 		"""Handle all types of errors that can occur during a step"""
 		error_msg = AgentError.format_error(error)
 		prefix = f'❌ Result failed {self.consecutive_failures + 1}/{self.max_failures} times:\n '
diff --git a/browser_use/browser/tests/playwright_test.py b/browser_use/browser/tests/playwright_test.py
index 1107841e0..4ac45a937 100644
--- a/browser_use/browser/tests/playwright_test.py
+++ b/browser_use/browser/tests/playwright_test.py
@@ -19,6 +19,8 @@ async def test_has_title(page: Page):
 
 	# Get all DOM content including all shadow roots recursively
 	start_time = time.time()
+	# wait for the page to load
+	await page.wait_for_load_state('load')
 	full_content = await dom_service._get_html_content()
 	# full_content = page.evaluate("""() => {
 	# 	function getAllContent(root) {
diff --git a/browser_use/dom/service.py b/browser_use/dom/service.py
index aaf7d6b1c..37349e57d 100644
--- a/browser_use/dom/service.py
+++ b/browser_use/dom/service.py
@@ -36,6 +36,7 @@ class DomService:
 		if with_shadow_roots:
 			full_content = await self.page.evaluate("""() => {
 				function getAllContent(root) {
+					if (!root) return '';
 					let content = root.innerHTML || '';
 					
 					// Get all elements with shadow roots