2
2
import json
3
3
import os
4
4
import shutil
5
- import tempfile
6
5
import signal
7
6
import sys
7
+ import tempfile
8
8
import time
9
9
from pathlib import Path
10
10
from typing import Any , Literal , Optional
@@ -47,7 +47,7 @@ class Stagehand:
47
47
48
48
# Dictionary to store one lock per session_id
49
49
_session_locks = {}
50
-
50
+
51
51
# Flag to track if cleanup has been called
52
52
_cleanup_called = False
53
53
@@ -122,15 +122,17 @@ def __init__(
122
122
self .wait_for_captcha_solves = self .config .wait_for_captcha_solves
123
123
self .system_prompt = self .config .system_prompt
124
124
self .verbose = self .config .verbose
125
-
125
+
126
126
# Smart environment detection
127
127
if self .config .env :
128
128
self .env = self .config .env .upper ()
129
129
else :
130
130
# 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
+ )
132
134
has_local_config = bool (self .config .local_browser_launch_options )
133
-
135
+
134
136
if has_local_config and not has_browserbase_config :
135
137
# Local browser options specified but no Browserbase config
136
138
self .env = "LOCAL"
@@ -140,7 +142,7 @@ def __init__(
140
142
else :
141
143
# Default to BROWSERBASE if Browserbase config is available
142
144
self .env = "BROWSERBASE"
143
-
145
+
144
146
self .local_browser_launch_options = (
145
147
self .config .local_browser_launch_options or {}
146
148
)
@@ -211,7 +213,7 @@ def __init__(
211
213
raise ValueError (
212
214
"browserbase_project_id is required for BROWSERBASE env with existing session_id (or set BROWSERBASE_PROJECT_ID in env)."
213
215
)
214
-
216
+
215
217
# Register signal handlers for graceful shutdown
216
218
self ._register_signal_handlers ()
217
219
@@ -242,16 +244,21 @@ def __init__(
242
244
243
245
def _register_signal_handlers (self ):
244
246
"""Register signal handlers for SIGINT and SIGTERM to ensure proper cleanup."""
247
+
245
248
def cleanup_handler (sig , frame ):
246
249
# Prevent multiple cleanup calls
247
250
if self .__class__ ._cleanup_called :
248
251
return
249
252
250
253
self .__class__ ._cleanup_called = True
251
254
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
+ )
253
258
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
+ )
255
262
256
263
try :
257
264
# Try to get the current event loop
@@ -275,9 +282,9 @@ def schedule_cleanup():
275
282
# Shield the task to prevent it from being cancelled
276
283
shielded = asyncio .shield (task )
277
284
# We don't need to await here since we're in call_soon_threadsafe
278
-
285
+
279
286
loop .call_soon_threadsafe (schedule_cleanup )
280
-
287
+
281
288
except Exception as e :
282
289
print (f"Error during signal cleanup: { str (e )} " )
283
290
sys .exit (1 )
0 commit comments