From d2e3de97d51aeecebfaad9850220c679e64d8fa7 Mon Sep 17 00:00:00 2001
From: magmueller <mamagnus00@gmail.com>
Date: Mon, 25 Nov 2024 16:56:15 +0100
Subject: [PATCH] Save cookies

---
 browser_use/agent/message_manager/service.py |  2 +-
 browser_use/browser/service.py               | 40 ++++++++++++++++++--
 browser_use/controller/service.py            |  6 ++-
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/browser_use/agent/message_manager/service.py b/browser_use/agent/message_manager/service.py
index e687d3038..f73905c1f 100644
--- a/browser_use/agent/message_manager/service.py
+++ b/browser_use/agent/message_manager/service.py
@@ -47,7 +47,7 @@ class MessageManager:
 			self.action_descriptions, current_date=datetime.now()
 		).get_system_message()
 		self._add_message_with_tokens(system_message)
-
+		self.system_prompt = system_message
 		task_message = HumanMessage(content=f'Your task is: {task}')
 		self._add_message_with_tokens(task_message)
 
diff --git a/browser_use/browser/service.py b/browser_use/browser/service.py
index da9d702c8..21e175f85 100644
--- a/browser_use/browser/service.py
+++ b/browser_use/browser/service.py
@@ -4,9 +4,12 @@ Playwright browser on steroids.
 
 import asyncio
 import base64
+import json
 import logging
+import os
 import time
 from dataclasses import dataclass
+from pathlib import Path
 
 from playwright.async_api import Browser as PlaywrightBrowser
 from playwright.async_api import BrowserContext, ElementHandle, Page, Playwright, async_playwright
@@ -34,9 +37,12 @@ class Browser:
 	MINIMUM_WAIT_TIME = 0.5
 	MAXIMUM_WAIT_TIME = 5
 
-	def __init__(self, headless: bool = False, keep_open: bool = False):
+	def __init__(
+		self, headless: bool = False, keep_open: bool = False, cookies_path: str | None = None
+	):
 		self.headless = headless
 		self.keep_open = keep_open
+		self.cookies_file = cookies_path
 
 		# Initialize these as None - they'll be set up when needed
 		self.session: BrowserSession | None = None
@@ -109,7 +115,7 @@ class Browser:
 			raise
 
 	async def _create_context(self, browser: PlaywrightBrowser):
-		"""Creates a new browser context with anti-detection measures."""
+		"""Creates a new browser context with anti-detection measures and loads cookies if available."""
 		context = await browser.new_context(
 			viewport={'width': 1280, 'height': 1024},
 			user_agent=(
@@ -119,6 +125,12 @@ class Browser:
 			java_script_enabled=True,
 		)
 
+		# Load cookies if they exist
+		if self.cookies_file and os.path.exists(self.cookies_file):
+			with open(self.cookies_file, 'r') as f:
+				cookies = json.load(f)
+				await context.add_cookies(cookies)
+
 		# Expose anti-detection scripts
 		await context.add_init_script(
 			"""
@@ -182,13 +194,22 @@ class Browser:
 
 	async def close(self, force: bool = False):
 		"""Close the browser instance"""
+
+		# check if already closed
+		if self.session is None:
+			return
+
+		if self.cookies_file:
+			await self.save_cookies()
+
 		if force and not self.keep_open:
 			session = await self.get_session()
 			await session.browser.close()
 			await session.playwright.stop()
+
 		else:
 			# Note: input() is blocking - consider an async alternative if needed
-			# input('Press Enter to close Browser...')
+			input('Press Enter to close Browser...')
 			self.keep_open = False
 			await self.close(force=True)
 
@@ -472,4 +493,15 @@ class Browser:
 			await self.get_xpath(index), timeout=2500, state='visible'
 		)
 
-	# endregion
+	async def save_cookies(self):
+		"""Save current cookies to file"""
+		if self.session and self.session.context and self.cookies_file:
+			try:
+				cookies = await self.session.context.cookies()
+				# maybe file is just a file name then i get
+				if os.path.dirname(self.cookies_file):
+					os.makedirs(os.path.dirname(self.cookies_file), exist_ok=True)
+				with open(self.cookies_file, 'w') as f:
+					json.dump(cookies, f)
+			except Exception as e:
+				logger.error(f'Failed to save cookies: {str(e)}')
diff --git a/browser_use/controller/service.py b/browser_use/controller/service.py
index fa472cd41..ed84448c2 100644
--- a/browser_use/controller/service.py
+++ b/browser_use/controller/service.py
@@ -22,8 +22,10 @@ logger = logging.getLogger(__name__)
 
 
 class Controller:
-	def __init__(self, headless: bool = False, keep_open: bool = False):
-		self.browser = Browser(headless=headless, keep_open=keep_open)
+	def __init__(
+		self, headless: bool = False, keep_open: bool = False, cookies_path: str | None = None
+	):
+		self.browser = Browser(headless=headless, keep_open=keep_open, cookies_path=cookies_path)
 		self.registry = Registry()
 		self._register_default_actions()