Implement Browserbase sessions in the Python browser factories

packages/sdk-python/README.md

@@ -30,6 +30,16 @@ asyncio.run(main())
 
 See [`examples`](examples) for action, extraction, observation, and custom LLM usage.
 
+For hosted browsers, import `browserbase` from `stagehand` and use
+`browserbase.launch(api_key=...)` or `browserbase.connect(api_key=..., session_id=...)` instead of
+`local_browser`. A launched session is released when you `close()` the browser handle, except with
+`keep_alive=True`: that handle's `close()` only disconnects, and the session keeps running until it
+is released out of band (Browserbase dashboard or API) or reaches its configured timeout. The
+Stagehand extension that `launch()` uploaded for that session is retained on the account for the
+same reason, so a `keep_alive=True` workflow that launches repeatedly accumulates extensions until
+they are deleted out of band. Sessions reached through `browserbase.connect()` are never released by
+`close()` — whoever created the session owns it.
+
 `Stagehand.act()`, `Stagehand.observe()`, and `Stagehand.extract()` use the active page by
 default. Pass `page=page` to target a specific SDK `Page`.
 

packages/sdk-python/pyproject.toml

@@ -4,7 +4,7 @@ version = "4.0.0"
 description = "Browserbase's SDK for building browser agents"
 readme = "README.md"
 requires-python = ">=3.11"
-dependencies = ["pydantic>=2.12,<3", "websockets>=16.1.1"]
+dependencies = ["browserbase>=1.15,<2", "pydantic>=2.12,<3", "websockets>=16.1.1"]
 license = "MIT"
 authors = [{ name = "Browserbase" }]
 keywords = ["ai", "browser", "automation", "web-scraping", "testing"]

packages/sdk-python/src/stagehand/browser.py

@@ -7,19 +7,28 @@
 import socket
 import sys
 import tempfile
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Literal, Protocol
-
-from ._generated.models import BrowserSessionMetadata
+from typing import Any, Literal, Protocol
+
+from ._generated.models import (
+    BrowserbaseBrowserSettings,
+    BrowserbaseRegion,
+    BrowserbaseSessionCreateParams,
+    BrowserSessionMetadata,
+    ProxyConfig,
+)
+from .browserbase_session import _create_browserbase_session_client
 from .cdp_client import CDPClient
 from .client_models import (
+    BrowserbaseConnectOptions,
     LocalBrowserConnectOptions,
     LocalBrowserLaunchOptions,
     LocalProxyConfig,
     LocalViewport,
 )
+from .extension_assets import extension_directory
 
 _WEBMCP_CHROME_FLAG = "--enable-features=WebMCPTesting,DevToolsWebMCPSupport"
 
@@ -152,6 +161,22 @@ def _release_browser(browser: StagehandBrowser) -> None:
     browser._claimed = False
 
 
+def _browser_session_metadata(
+    session_id: str,
+    region: BrowserbaseRegion | None,
+) -> BrowserSessionMetadata:
+    # Leave region unset (not None) so the wire payload omits it entirely; the
+    # worker schema rejects "region": null.
+    return BrowserSessionMetadata.model_validate({
+        name: value
+        for name, value in (
+            ("session_id", session_id),
+            ("region", region),
+        )
+        if value is not None
+    })
+
+
 class _BrowserConnectionSource(Protocol):
     cdp_url: str
     keep_alive: bool
@@ -188,6 +213,7 @@ async def _connect_browser(
     source: _BrowserConnectionSource,
     extension_dir: str | None = None,
     extension_id: str | None = None,
+    preloaded_extension: bool = False,
     connect_timeout_ms: int | None = None,
     after_connect: Callable[[CDPClient], Awaitable[None]] | None = None,
     worker_init_metadata: _WorkerInitMetadata,
@@ -200,6 +226,7 @@ async def _connect_browser(
             cdp_url=source.cdp_url,
             extension_dir=extension_dir,
             extension_id=extension_id,
+            preloaded_extension=preloaded_extension,
             service_worker_url_includes="service-worker.js",
             discovery_timeout_ms=timeout_ms,
             command_timeout_ms=_COMMAND_TIMEOUT_MS,
@@ -252,13 +279,6 @@ async def close() -> None:
     )
 
 
-def _resolve_extension_dir() -> Path:
-    extension_dir = Path(__file__).with_name("_extension")
-    if not (extension_dir / "manifest.json").is_file():
-        extension_dir = Path(__file__).resolve().parents[3] / "server" / "dist"
-    return extension_dir
-
-
 class LocalBrowser:
     async def launch(
         self,
@@ -384,7 +404,7 @@ async def configure_downloads(cdp_client: CDPClient) -> None:
             provider="local",
             origin="launched",
             source=source,
-            extension_dir=str(_resolve_extension_dir()),
+            extension_dir=str(extension_directory()),
             connect_timeout_ms=options.connect_timeout_ms,
             after_connect=(
                 configure_downloads
@@ -410,7 +430,7 @@ async def connect(
             )
             if value is not None
         })
-        extension_dir = None if options.extension_id is not None else str(_resolve_extension_dir())
+        extension_dir = None if options.extension_id is not None else str(extension_directory())
         return await _connect_browser(
             provider="local",
             origin="connected",
@@ -432,11 +452,91 @@ async def close(self) -> None:
 
 
 class BrowserbaseBrowser:
-    async def launch(self, **_kwargs: object) -> StagehandBrowser:
-        raise NotImplementedError("Browserbase sessions are not implemented yet")
+    async def launch(
+        self,
+        *,
+        api_key: str,
+        browser_settings: BrowserbaseBrowserSettings | None = None,
+        extension_id: str | None = None,
+        keep_alive: bool | None = None,
+        proxies: bool | list[ProxyConfig] | None = None,
+        region: BrowserbaseRegion | None = None,
+        timeout: float | None = None,
+        user_metadata: Mapping[str, Any] | None = None,
+    ) -> StagehandBrowser:
+        if not api_key:
+            raise ValueError("api_key must not be empty")
+        if extension_id is not None and not extension_id.strip():
+            raise ValueError("extension_id must not be empty")
+        if (
+            browser_settings is not None
+            and browser_settings.extension_id is not None
+            and not browser_settings.extension_id.strip()
+        ):
+            raise ValueError("browser_settings.extension_id must not be empty")
+        options = BrowserbaseSessionCreateParams.model_validate({
+            name: value
+            for name, value in (
+                ("browser_settings", browser_settings),
+                ("extension_id", extension_id),
+                ("keep_alive", keep_alive),
+                ("proxies", proxies),
+                ("region", region),
+                ("timeout", timeout),
+                ("user_metadata", dict(user_metadata) if user_metadata is not None else None),
+            )
+            if value is not None
+        })
+        session = await _create_browserbase_session_client(api_key).create_session(options)
+        source = ResolvedBrowserSource(
+            cdp_url=session.cdp_url,
+            keep_alive=options.keep_alive or False,
+            _close_callback=session.close,
+        )
+        return await _connect_browser(
+            provider="browserbase",
+            origin="launched",
+            source=source,
+            preloaded_extension=True,
+            worker_init_metadata=_WorkerInitMetadata(
+                api_key=api_key,
+                browser=_browser_session_metadata(session.session_id, options.region),
+            ),
+        )
 
-    async def connect(self, **_kwargs: object) -> StagehandBrowser:
-        raise NotImplementedError("Browserbase sessions are not implemented yet")
+    async def connect(
+        self,
+        *,
+        api_key: str,
+        session_id: str,
+        connect_timeout_ms: int | None = None,
+        extension_id: str | None = None,
+    ) -> StagehandBrowser:
+        options = BrowserbaseConnectOptions.model_validate({
+            name: value
+            for name, value in (
+                ("api_key", api_key),
+                ("session_id", session_id),
+                ("connect_timeout_ms", connect_timeout_ms),
+                ("extension_id", extension_id),
+            )
+            if value is not None
+        })
+        connection = await _create_browserbase_session_client(options.api_key).connect_session(
+            options.session_id
+        )
+        return await _connect_browser(
+            provider="browserbase",
+            origin="connected",
+            source=_ConnectedBrowserSource(connection.cdp_url),
+            extension_id=options.extension_id,
+            preloaded_extension=options.extension_id is None,
+            connect_timeout_ms=options.connect_timeout_ms,
+            worker_init_metadata=_WorkerInitMetadata(
+                api_key=options.api_key,
+                browser=_browser_session_metadata(connection.session_id, connection.region),
+            ),
+        )
 
 
 local_browser = LocalBrowser()