Skip to content

Commit 8ca7751

Browse files
format
1 parent 37d5510 commit 8ca7751

File tree

1 file changed

+18
-11
lines changed

1 file changed

+18
-11
lines changed

stagehand/client.py

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
import json
33
import os
44
import shutil
5-
import tempfile
65
import signal
76
import sys
7+
import tempfile
88
import time
99
from pathlib import Path
1010
from typing import Any, Literal, Optional
@@ -47,7 +47,7 @@ class Stagehand:
4747

4848
# Dictionary to store one lock per session_id
4949
_session_locks = {}
50-
50+
5151
# Flag to track if cleanup has been called
5252
_cleanup_called = False
5353

@@ -122,15 +122,17 @@ def __init__(
122122
self.wait_for_captcha_solves = self.config.wait_for_captcha_solves
123123
self.system_prompt = self.config.system_prompt
124124
self.verbose = self.config.verbose
125-
125+
126126
# Smart environment detection
127127
if self.config.env:
128128
self.env = self.config.env.upper()
129129
else:
130130
# Auto-detect environment based on available configuration
131-
has_browserbase_config = bool(self.browserbase_api_key and self.browserbase_project_id)
131+
has_browserbase_config = bool(
132+
self.browserbase_api_key and self.browserbase_project_id
133+
)
132134
has_local_config = bool(self.config.local_browser_launch_options)
133-
135+
134136
if has_local_config and not has_browserbase_config:
135137
# Local browser options specified but no Browserbase config
136138
self.env = "LOCAL"
@@ -140,7 +142,7 @@ def __init__(
140142
else:
141143
# Default to BROWSERBASE if Browserbase config is available
142144
self.env = "BROWSERBASE"
143-
145+
144146
self.local_browser_launch_options = (
145147
self.config.local_browser_launch_options or {}
146148
)
@@ -211,7 +213,7 @@ def __init__(
211213
raise ValueError(
212214
"browserbase_project_id is required for BROWSERBASE env with existing session_id (or set BROWSERBASE_PROJECT_ID in env)."
213215
)
214-
216+
215217
# Register signal handlers for graceful shutdown
216218
self._register_signal_handlers()
217219

@@ -242,16 +244,21 @@ def __init__(
242244

243245
def _register_signal_handlers(self):
244246
"""Register signal handlers for SIGINT and SIGTERM to ensure proper cleanup."""
247+
245248
def cleanup_handler(sig, frame):
246249
# Prevent multiple cleanup calls
247250
if self.__class__._cleanup_called:
248251
return
249252

250253
self.__class__._cleanup_called = True
251254
if self.env == "BROWSERBASE":
252-
print(f"\n[{signal.Signals(sig).name}] received. Ending Browserbase session...")
255+
print(
256+
f"\n[{signal.Signals(sig).name}] received. Ending Browserbase session..."
257+
)
253258
else:
254-
print(f"\n[{signal.Signals(sig).name}] received. Cleaning up Stagehand resources...")
259+
print(
260+
f"\n[{signal.Signals(sig).name}] received. Cleaning up Stagehand resources..."
261+
)
255262

256263
try:
257264
# Try to get the current event loop
@@ -275,9 +282,9 @@ def schedule_cleanup():
275282
# Shield the task to prevent it from being cancelled
276283
shielded = asyncio.shield(task)
277284
# We don't need to await here since we're in call_soon_threadsafe
278-
285+
279286
loop.call_soon_threadsafe(schedule_cleanup)
280-
287+
281288
except Exception as e:
282289
print(f"Error during signal cleanup: {str(e)}")
283290
sys.exit(1)

0 commit comments

Comments
 (0)